From 859529ccbc1f3c0bf9daace4a9c92a84e60f439f Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 16 Oct 2026 15:24:40 +0000 Subject: [PATCH 01/58] fix(realtime)!: keep turn_detection.type when a transcript guardrail disables auto-response (#44939) * fix(realtime): keep turn_detection.type when a transcript guardrail disables auto-response * test(realtime): type the GA session.update helper and mark its locals Final * test(realtime): expect typed turn_detection on the voice guardrail session updates * test(realtime): accept either gate path for the client's first voice session.update * fix(gemini): drop VAD tuning when the guardrail turns auto-response off * fix(gemini): check create_response without a redundant isinstance * fix(vertex_ai): keep the client's VAD tuning on Vertex when auto-response is off * fix(realtime): keep the transcript-guardrail gate typeless on xAI so it does not auto-answer * fix(realtime): replace any client turn_detection with the untyped gate on xAI xAI answers on its own under a typed server_vad even with create_response off, so a client's typed turn_detection must not reach xAI while the transcript guardrail gates the turn. * fix(realtime): keep a client's early turn_detection and hide the guardrail cancel's error A session.update sent before session.created reached OpenAI ahead of the proxy's own gate, so the gate replaced the client's turn_detection. The proxy now re-sends the client's gated turn_detection after the gate when the two differ. The guardrail's response.cancel now carries a litellm_guardrail_ event_id, and the error OpenAI returns for it (response_cancel_not_active when nothing is answering) no longer reaches the client. * fix(realtime): send the guardrail gate before a client's early session.update and let Muse and Chirp fall back to server_vad under drop_params * fix(realtime): hide only the cancel errors that echo the guardrail's own response.cancel ids Track the event ids of the guardrail cancels this session sent and drop a backend error only when its error.event_id is one of them, so a client's own event id with the same prefix still gets its error. Cover the scripted response.cancel reply in the integration upstream, update the stale list-mode test to the typed gate, and mark the xAI cancel error cell as an expected failure, since xAI replies with its own event id. * test(integration support): deliver the disconnect a real server sends after the client goes away The scripted upstream now keeps reading frames after a reply hits a closed socket, so the unit test's fake ASGI receive has to end with the websocket.disconnect message a real server delivers. * fix(realtime): delete the blocked transcript item before the guardrail's refusal turn OpenAI GA answered the blocked question in some refusal turns because the blocked user item stayed in the conversation. The raw OpenAI-protocol path now sends conversation.item.delete for it right after the cancel. Provider config backends have no such event, so that path no longer passes an item id. Also bound the guardrail integration file by the shared graceful-stop timeout and make the xAI cancel-error xfail require an AssertionError. * fix(realtime): keep the blocked item on xAI, which cannot delete an audio input item * test(realtime): assert the xAI voice block sends no conversation.item.delete upstream --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../litellm_core_utils/realtime_streaming.py | 228 ++++++--- .../realtime/transcription_protocol.py | 10 + .../llms/gemini/realtime/transformation.py | 11 +- litellm/llms/meta/realtime/transformation.py | 17 +- .../realtime_transformation.py | 18 +- .../llms/vertex_ai/realtime/transformation.py | 8 +- litellm/llms/xai/realtime/transformation.py | 13 + litellm/realtime_api/main.py | 4 +- litellm/utils.py | 3 +- tests/integration/_support/upstream.py | 198 +++++++- .../cost_calculation/cost_tracking_case.py | 2 + ..._transcription_guardrail_session_update.py | 473 ++++++++++++++++-- .../test_realtime_unguarded_voice_turns.py | 14 +- .../unit/integration_support/test_upstream.py | 6 +- .../test_realtime_streaming.py | 320 +++++++++++- .../test_gemini_realtime_transformation.py | 36 +- .../test_meta_realtime_transformation.py | 18 + .../test_vertex_ai_realtime_transformation.py | 21 + .../test_vertex_ai_realtime_transformation.py | 30 ++ 19 files changed, 1252 insertions(+), 178 deletions(-) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index c9b96c0a4c2..e4eb9acc37d 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -1,16 +1,18 @@ import asyncio import json import traceback -from collections.abc import Coroutine, Mapping, Sequence +from collections.abc import Callable, Coroutine, Mapping, Sequence from dataclasses import dataclass from enum import Enum, auto -from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypedDict, cast +from types import MappingProxyType +from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypeAlias, TypedDict, cast from pydantic import TypeAdapter from typing_extensions import ReadOnly import litellm from litellm._logging import redact_internal_details_from_client_message, verbose_logger +from litellm._uuid import uuid from litellm.constants import ( REALTIME_METERED_USAGE_KEY, REALTIME_SESSION_FAILURE_LOGGED_KEY, @@ -116,10 +118,87 @@ def _decode_json_object(payload: str) -> Mapping[str, object]: return json.loads(payload) +_GUARDRAIL_TURN_DETECTION_GATE: Final[Mapping[str, object]] = MappingProxyType( + {"type": "server_vad", "create_response": False} +) +_EMPTY_SESSION: Final[Mapping[str, object]] = MappingProxyType({}) +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) +_GUARDRAIL_EVENT_ID_PREFIX: Final = "litellm_guardrail_" + + +def _json_object(value: object) -> Mapping[str, object] | None: + return _JSON_OBJECT.validate_python(value) if isinstance(value, Mapping) else None + + +def _mapping_field(container: Mapping[str, object], key: str) -> Mapping[str, object] | None: + return _json_object(container.get(key)) + + +def _guardrail_event_id() -> str: + return f"{_GUARDRAIL_EVENT_ID_PREFIX}{uuid.uuid4().hex}" + + +def _error_event_id(event: Mapping[str, object]) -> str | None: + error: Final = _mapping_field(event, "error") + caused_by: Final = None if error is None else error.get("event_id") + return caused_by if isinstance(caused_by, str) else None + + +_TurnDetectionGate: TypeAlias = Callable[[object], Mapping[str, object]] + + +def _turn_detection_with_auto_response_off(turn_detection: object) -> Mapping[str, object]: + client_turn_detection: Final = _json_object(turn_detection) + if client_turn_detection is None: + return dict(_GUARDRAIL_TURN_DETECTION_GATE) + return {**client_turn_detection, "create_response": False} + + +def _nested_turn_detection_part(session: Mapping[str, object], gate: _TurnDetectionGate) -> Mapping[str, object]: + audio: Final = _mapping_field(session, "audio") + if audio is None: + return {} + audio_input: Final = _mapping_field(audio, "input") + if audio_input is None or "turn_detection" not in audio_input: + return {} + gated_input: Final = { + **audio_input, + "turn_detection": gate(audio_input["turn_detection"]), + } + return {"audio": {**audio, "input": gated_input}} + + +def _with_create_response_disabled( + session: Mapping[str, object], gate: _TurnDetectionGate +) -> Mapping[str, object] | None: + """``session`` with every turn_detection the client sent unable to auto-respond, or None when it sent none.""" + flat_part: Final = {"turn_detection": gate(session["turn_detection"])} if "turn_detection" in session else {} + nested_part: Final = _nested_turn_detection_part(session, gate) + if not flat_part and not nested_part: + return None + return {**session, **flat_part, **nested_part} + + +def _gate_session_update( + session: object, gate: _TurnDetectionGate, inject_gate_when_absent: bool +) -> Mapping[str, object] | None: + """The session to forward under a transcript guardrail, or None when the frame needs no change.""" + client_session: Final = _json_object(session) + if client_session is None: + return None + gated: Final = _with_create_response_disabled(client_session, gate) + if gated is not None or not inject_gate_when_absent: + return gated + return {**client_session, "turn_detection": gate(None)} + + class RealtimeEventNormalizer(Protocol): def should_drop(self, event: object) -> bool: ... def normalize(self, event: dict) -> dict: ... def patch_outgoing_session(self, session: dict) -> dict: ... + def gate_turn_detection(self, turn_detection: object, /) -> Mapping[str, object]: ... + @property + def deletes_conversation_items(self) -> bool: ... class RealtimeUsageMeter(Protocol): @@ -197,6 +276,8 @@ class RealTimeStreaming: # Track whether we have already sent the guardrail turn-detection update # that disables provider auto-response for transcription guardrails. self._guardrail_turn_detection_update_sent: bool = False + self._session_gate_sent: bool = False + self._guardrail_cancel_event_ids: frozenset[str] = frozenset() # Deferred Gemini Live setup: Pipecat may stream audio before session.update. # Buffer client audio until the backend acknowledges setup (setupComplete). self._backend_setup_complete: bool = provider_config is None or provider_config.requires_session_configuration() @@ -213,6 +294,12 @@ class RealTimeStreaming: self._is_transcription_session: bool = force_transcription_model is not None # Optional per-provider GA event normalizer (e.g. XAIRealtimeNormalizer). self._event_normalizer = event_normalizer + self._gate_turn_detection: Final[_TurnDetectionGate] = ( + _turn_detection_with_auto_response_off if event_normalizer is None else event_normalizer.gate_turn_detection + ) + self._deletes_conversation_items: Final = ( + event_normalizer is None or event_normalizer.deletes_conversation_items + ) self._usage_meter = usage_meter @property @@ -734,20 +821,37 @@ class RealTimeStreaming: except (json.JSONDecodeError, TypeError): return + def _make_turn_detection_update(self, turn_detection: Mapping[str, object]) -> dict[str, object]: + session: Final[dict[str, object]] = ( + {"turn_detection": turn_detection} + if self._backend_uses_beta_protocol + else {"type": "realtime", "audio": {"input": {"turn_detection": turn_detection}}} + ) + return {"type": "session.update", "session": session} + def _make_disable_auto_response_message(self) -> str: """Return a session.update that disables VAD auto-response.""" - turn_detection: Final[dict[str, str | bool]] = { - "type": "server_vad", - "create_response": False, - } - if self._backend_uses_beta_protocol: - session: dict[str, object] = {"turn_detection": turn_detection} - else: - session = { - "type": "realtime", - "audio": {"input": {"turn_detection": turn_detection}}, - } - return json.dumps({"type": "session.update", "session": session}) + return json.dumps(self._make_turn_detection_update(self._gate_turn_detection(None))) + + async def _send_guardrail_cancel(self) -> None: + event_id: Final = _guardrail_event_id() + self._guardrail_cancel_event_ids = self._guardrail_cancel_event_ids | {event_id} + await self._send_to_backend(json.dumps({"type": "response.cancel", "event_id": event_id})) + + async def _withdraw_blocked_turn(self, item_id: str | None) -> None: + await self._send_guardrail_cancel() + if item_id is not None and self._deletes_conversation_items: + await self._send_to_backend(json.dumps({"type": "conversation.item.delete", "item_id": item_id})) + + def _is_guardrail_cancel_error(self, event: Mapping[str, object]) -> bool: + return _error_event_id(event) in self._guardrail_cancel_event_ids + + async def _send_session_gate_once(self) -> None: + """Gate the backend's default turn_detection ahead of the first client session.update, which overrides it.""" + if self._session_gate_sent or not self._should_disable_vad_auto_response(): + return + self._session_gate_sent = True + await self._send_to_backend(self._make_disable_auto_response_message()) async def _maybe_send_guardrail_turn_detection_update(self) -> None: """Disable provider auto-response once when transcription guardrails are enabled.""" @@ -921,7 +1025,7 @@ class RealTimeStreaming: if pre_block_backend_message is not None: await self._send_to_backend(pre_block_backend_message) if not self._is_transcription_session: - await self._send_to_backend(json.dumps({"type": "response.cancel"})) + await self._withdraw_blocked_turn(item_id) # Send the policy violation hint (shows as small gray status text in UI). await self.websocket.send_text( json.dumps( @@ -1051,10 +1155,7 @@ class RealTimeStreaming: self.store_message(event_str) self._capture_transcription_usage(event) await self._send_event_to_client(event, event_str) - blocked = await self.run_realtime_guardrails( - cast(str, transcript), - item_id=cast(str | None, event.get("item_id")), - ) + blocked = await self.run_realtime_guardrails(cast(str, transcript)) if not blocked and self._should_disable_vad_auto_response(): await self._send_to_backend(json.dumps({"type": "response.create"})) continue @@ -1086,7 +1187,7 @@ class RealTimeStreaming: if event_type == "session.created" and self._should_disable_vad_auto_response(): self.store_message(event_obj) await self.websocket.send_text(self._event_to_client_json(event_obj)) - await self._send_to_backend(self._make_disable_auto_response_message()) + await self._send_session_gate_once() return True if event_type == "conversation.item.input_audio_transcription.completed": @@ -1140,6 +1241,10 @@ class RealTimeStreaming: if self._usage_meter is not None: self._usage_meter.observe_backend_event(event) + if event.get("type") == "error" and self._is_guardrail_cancel_error(event): + verbose_logger.debug("Realtime: not forwarding an error caused by a guardrail event: %s", event) + continue + if await self._handle_raw_backend_message(event, raw_response): continue @@ -1472,19 +1577,21 @@ class RealTimeStreaming: and self.session_configuration_request is None and not self._guardrail_turn_detection_update_sent and self._should_disable_vad_auto_response() - ): - session: Mapping[str, object] | None = msg_obj.setdefault("session", {}) - if isinstance(session, dict): - existing_td = session.get("turn_detection") - if not isinstance(existing_td, dict): - existing_td = {} - existing_td["create_response"] = False - session["turn_detection"] = existing_td - message = json.dumps(msg_obj) - guardrail_turn_detection_injected = True - verbose_logger.debug( - "Injected turn_detection into first session.update for audio transcription guardrails" + and ( + gated_first := _gate_session_update( + client_event.get("session", _EMPTY_SESSION), + self._gate_turn_detection, + inject_gate_when_absent=True, ) + ) + is not None + ): + msg_obj["session"] = gated_first + message = json.dumps(msg_obj) + guardrail_turn_detection_injected = True + verbose_logger.debug( + "Injected turn_detection into first session.update for audio transcription guardrails" + ) ## GUARDRAIL: Force ``create_response`` to False in any # client-provided ``turn_detection`` so a later @@ -1492,52 +1599,25 @@ class RealTimeStreaming: # and bypass the transcription guardrail after the # initial disable. Covers both the flat beta key and the # nested GA ``audio.input.turn_detection`` shape, since - # the GA remap below also accepts either form. Skipped - # when the injection block above already ran for this - # message, to avoid redundant double-serialization. + # the GA remap below also accepts either form. A later + # update with no turn_detection is forwarded as is: OpenAI + # keeps the earlier disable for an unset field, and xAI + # falls back to its manual default, which never auto-responds. if ( msg_type == "session.update" and not guardrail_turn_detection_injected and self._should_disable_vad_auto_response() + and ( + gated_later := _gate_session_update( + client_event.get("session"), + self._gate_turn_detection, + inject_gate_when_absent=False, + ) + ) + is not None ): - session = client_event.get("session") - if isinstance(session, dict): - td_overridden = False - flat_td = session.get("turn_detection") - flat_td_present = flat_td is not None - if flat_td_present: - if not isinstance(flat_td, dict): - flat_td = {} - if flat_td.get("create_response") is not False: - flat_td["create_response"] = False - session["turn_detection"] = flat_td - td_overridden = True - nested_td_present = False - audio = session.get("audio") - if isinstance(audio, dict): - audio_input = audio.get("input") - if isinstance(audio_input, dict): - nested_td = audio_input.get("turn_detection") - if nested_td is not None: - nested_td_present = True - if not isinstance(nested_td, dict): - nested_td = {} - if nested_td.get("create_response") is not False: - nested_td["create_response"] = False - audio_input["turn_detection"] = nested_td - td_overridden = True - # Symmetric with the first-update injection block: - # if the client omitted turn_detection entirely on - # a subsequent session.update, still inject the - # ``create_response: False`` override so the - # transcription guardrail cannot be re-enabled by - # any downstream merge that drops the original - # disable. - if not flat_td_present and not nested_td_present: - session["turn_detection"] = {"create_response": False} - td_overridden = True - if td_overridden: - message = json.dumps(msg_obj) + msg_obj["session"] = gated_later + message = json.dumps(msg_obj) # GA compatibility: remap beta-style session fields only when # the upstream is in GA mode. Beta upstreams expect the flat @@ -1590,6 +1670,8 @@ class RealTimeStreaming: # would permanently disable the injection if ``_send_to_backend`` # raised — neither this loop nor # ``_maybe_send_guardrail_turn_detection_update`` would retry. + if msg_type == "session.update" and self.provider_config is None: + await self._send_session_gate_once() sent = await self._send_to_backend(message) if guardrail_turn_detection_injected and sent: self._guardrail_turn_detection_update_sent = True diff --git a/litellm/llms/base_llm/realtime/transcription_protocol.py b/litellm/llms/base_llm/realtime/transcription_protocol.py index c3264911d46..b4514fd1da2 100644 --- a/litellm/llms/base_llm/realtime/transcription_protocol.py +++ b/litellm/llms/base_llm/realtime/transcription_protocol.py @@ -102,6 +102,16 @@ def json_integer( return value +def uses_server_vad( + update: TranscriptionSessionUpdate, backend: str, drop_params: bool, error: ProtocolErrorType +) -> bool: + if update.turn_detection_disabled: + return False + if update.turn_detection_type in (None, "server_vad") or drop_params: + return True + raise error(f"{backend} supports server_vad turn detection or null. Set drop_params to use server_vad instead") + + def new_event_id() -> str: return f"event_{uuid.uuid4().hex}" diff --git a/litellm/llms/gemini/realtime/transformation.py b/litellm/llms/gemini/realtime/transformation.py index 7c04ebf311f..f4e2941670a 100644 --- a/litellm/llms/gemini/realtime/transformation.py +++ b/litellm/llms/gemini/realtime/transformation.py @@ -161,6 +161,10 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): """Google AI Studio Gemini 3.5+ accepts ``id`` on functionResponses; Vertex AI rejects it.""" return True + def _rejects_vad_tuning_with_detection_off(self) -> bool: + """Google AI Studio closes the socket (1007) on padding or silence tuning next to ``disabled: true``.""" + return True + @staticmethod def _usage_detail_alias(details: Mapping[str, int | None] | None, defaults: dict[str, int]) -> dict[str, int]: if not isinstance(details, dict): @@ -279,10 +283,15 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): OpenAI ``semantic_vad`` has no Gemini Live equivalent — return an empty dict so callers omit ``realtimeInputConfig`` (mapping it with - ``disabled: true`` breaks native-audio sessions). + ``disabled: true`` breaks native-audio sessions). An explicit + ``create_response: false`` maps to detection off alone where the backend + rejects tuning next to ``disabled: true``. """ if isinstance(value, dict) and value.get("type") == "semantic_vad" and "create_response" not in value: return AutomaticActivityDetection() + auto_response_off: Final = "create_response" in value and value["create_response"] is False + if auto_response_off and self._rejects_vad_tuning_with_detection_off(): + return AutomaticActivityDetection(disabled=True) automatic_activity_dection: Final = AutomaticActivityDetection() if "create_response" in value and isinstance(value["create_response"], bool): diff --git a/litellm/llms/meta/realtime/transformation.py b/litellm/llms/meta/realtime/transformation.py index 46231e5e980..892d7ff66de 100644 --- a/litellm/llms/meta/realtime/transformation.py +++ b/litellm/llms/meta/realtime/transformation.py @@ -26,6 +26,7 @@ from litellm.llms.base_llm.realtime.transcription_protocol import ( speech_event, transcription_session, transcription_session_created_event, + uses_server_vad, ) from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig from litellm.secret_managers.main import get_secret_str @@ -230,15 +231,7 @@ def _parse_sample_rate(update: TranscriptionSessionUpdate) -> MuseSampleRate: return 16_000 if rate == 16_000 else 24_000 -def _parse_mode(update: TranscriptionSessionUpdate) -> MuseMode: - if update.turn_detection_disabled: - return "PUSH_TO_TALK" - if update.turn_detection_type not in (None, "server_vad"): - raise MuseProtocolError("Muse Voice supports server_vad turn detection or null") - return "ENDPOINTING" - - -def parse_session_update(payload: str, expected_model: str) -> MuseSessionConfig: +def parse_session_update(payload: str, expected_model: str, *, drop_params: bool = False) -> MuseSessionConfig: update: Final = parse_transcription_session_update(payload, MuseProtocolError) if update.session_type not in (None, "transcription", "realtime"): raise MuseProtocolError("Muse Voice supports transcription sessions only") @@ -253,7 +246,7 @@ def parse_session_update(payload: str, expected_model: str) -> MuseSessionConfig raise MuseProtocolError("realtime session model cannot be changed") return MuseSessionConfig( model=normalized_model, - mode=_parse_mode(update), + mode="ENDPOINTING" if uses_server_vad(update, "Muse Voice", drop_params, MuseProtocolError) else "PUSH_TO_TALK", sample_rate=_parse_sample_rate(update), language_bias=() if update.language is None else (normalize_language(update.language),), ) @@ -441,9 +434,11 @@ class MetaRealtimeConfig(BaseRealtimeConfig): *, monotonic: Callable[[], float] = time.monotonic, sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, + drop_params: bool = False, ) -> None: self._monotonic: Final = monotonic self._sleep: Final = sleep + self._drop_params: Final = drop_params self._transformer: Final = MuseEventTransformer() self._access_token: str | None = None self._config: MuseSessionConfig | None = None @@ -553,7 +548,7 @@ class MetaRealtimeConfig(BaseRealtimeConfig): access_token: Final = self._access_token if access_token is None: raise MuseProtocolError("Meta API key was not validated before the session was configured") - config: Final = parse_session_update(message, model) + config: Final = parse_session_update(message, model, drop_params=self._drop_params) self._config = config self._transformer.configure(config) return (json.dumps(config.handshake(access_token), separators=(",", ":")),) diff --git a/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py b/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py index 1d4a974fc54..5015678c467 100644 --- a/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py +++ b/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py @@ -13,7 +13,6 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.llms.base_llm.realtime.transcription_protocol import ( RealtimeTranscriptionProtocolError, TranscriptionAudioFormat, - TranscriptionSessionUpdate, completed_event, decode_pcm16_append, delta_event, @@ -23,6 +22,7 @@ from litellm.llms.base_llm.realtime.transcription_protocol import ( speech_event, transcription_session, transcription_session_created_event, + uses_server_vad, ) from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig, RealtimeBackend from litellm.llms.vertex_ai.audio_transcription.transformation import ( @@ -125,7 +125,7 @@ def default_session_config(model: str) -> ChirpSessionConfig: ) -def parse_chirp_session_update(payload: str, expected_model: str) -> ChirpSessionConfig: +def parse_chirp_session_update(payload: str, expected_model: str, *, drop_params: bool = False) -> ChirpSessionConfig: update: Final = parse_transcription_session_update(payload, ChirpProtocolError) if update.session_type not in (None, "transcription", "realtime"): raise ChirpProtocolError("Speech-to-Text streaming supports transcription sessions only") @@ -141,7 +141,7 @@ def parse_chirp_session_update(payload: str, expected_model: str) -> ChirpSessio model=model, language=None if update.language is None else normalize_transcription_language_to_bcp47(update.language), sample_rate=_parse_sample_rate(update.audio_format), - server_vad=_parse_server_vad(update), + server_vad=uses_server_vad(update, "Speech-to-Text streaming", drop_params, ChirpProtocolError), ) @@ -161,14 +161,6 @@ def _parse_sample_rate(audio_format: TranscriptionAudioFormat | None) -> int: return rate -def _parse_server_vad(update: TranscriptionSessionUpdate) -> bool: - if update.turn_detection_disabled: - return False - if update.turn_detection_type not in (None, "server_vad"): - raise ChirpProtocolError("Speech-to-Text streaming supports server_vad turn detection or null") - return True - - def session_created_event(config: ChirpSessionConfig, session_id: str) -> OpenAIRealtimeTranscriptionSessionCreated: return transcription_session_created_event(config.openai_session(session_id)) @@ -331,11 +323,13 @@ class VertexChirpRealtimeConfig(BaseRealtimeConfig): project: str, location: str | None, backend_factory: Callable[[SpeechStreamingTarget], RealtimeBackend] = _default_backend_factory, + drop_params: bool = False, ) -> None: self._resolve_access_token: Final = resolve_access_token self._project: Final = validate_vertex_transcription_project_id(project) self._location: Final = validate_vertex_transcription_location(location, DEFAULT_SPEECH_TO_TEXT_LOCATION) self._backend_factory: Final = backend_factory + self._drop_params: Final = drop_params self._transformer: Final = ChirpEventTransformer() self._config: ChirpSessionConfig | None = None self._session_id: str | None = None @@ -423,7 +417,7 @@ class VertexChirpRealtimeConfig(BaseRealtimeConfig): if self._config is not None: verbose_logger.debug("Speech-to-Text streaming: ignoring session.update after the stream was configured") return () - config: Final = parse_chirp_session_update(message, model) + config: Final = parse_chirp_session_update(message, model, drop_params=self._drop_params) self._config = config self._transformer.configure(config, self._session_id or f"sess_{uuid.uuid4().hex}") return (config.configure_command(),) diff --git a/litellm/llms/vertex_ai/realtime/transformation.py b/litellm/llms/vertex_ai/realtime/transformation.py index c14dd535145..55f1cf58cbe 100644 --- a/litellm/llms/vertex_ai/realtime/transformation.py +++ b/litellm/llms/vertex_ai/realtime/transformation.py @@ -47,6 +47,9 @@ class VertexAIRealtimeConfig(GeminiRealtimeConfig): def _include_function_response_id(self) -> bool: return False + def _rejects_vad_tuning_with_detection_off(self) -> bool: + return False + # ------------------------------------------------------------------ # URL # ------------------------------------------------------------------ @@ -248,9 +251,12 @@ def vertex_realtime_config( resolve_access_token: Callable[[], Awaitable[str]], project: str, location: str | None, + drop_params: bool, ) -> VertexAIRealtimeConfig | VertexChirpRealtimeConfig: if is_vertex_speech_to_text_model(model): - return VertexChirpRealtimeConfig(resolve_access_token=resolve_access_token, project=project, location=location) + return VertexChirpRealtimeConfig( + resolve_access_token=resolve_access_token, project=project, location=location, drop_params=drop_params + ) return VertexAIRealtimeConfig( access_token=access_token, project=project, diff --git a/litellm/llms/xai/realtime/transformation.py b/litellm/llms/xai/realtime/transformation.py index 5efe125ee60..ffbf814788b 100644 --- a/litellm/llms/xai/realtime/transformation.py +++ b/litellm/llms/xai/realtime/transformation.py @@ -16,6 +16,7 @@ construction time (see ``handler.py``) so all normalization is isolated here and ``RealTimeStreaming`` stays provider-agnostic. """ +from collections.abc import Mapping from typing import Final @@ -77,6 +78,18 @@ class XAIRealtimeNormalizer: event = self._normalize_response_usage_event(event, event_type) return event + def gate_turn_detection(self, _turn_detection: object, /) -> Mapping[str, object]: + """xAI auto-answers under ``server_vad`` even with ``create_response`` off, so the gate is always untyped.""" + return {"create_response": False} + + @property + def deletes_conversation_items(self) -> bool: + """xAI answers ``conversation.item.delete`` with ``Item not found`` until it commits the audio item. + + The guardrail blocks on an interim transcript, before that commit. + """ + return False + def patch_outgoing_session(self, session: dict) -> dict: """Patch a client ``session.update`` payload before forwarding to xAI. diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index 367d5be7bfc..aeec0bf67c0 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -500,11 +500,13 @@ async def _arealtime( ) return + drop_params: Final = litellm.drop_params is True or normalize_drop_params(kwargs.get("drop_params")) is True provider_config: BaseRealtimeConfig | None = None if _custom_llm_provider in LlmProviders._member_map_.values(): provider_config = ProviderConfigManager.get_provider_realtime_config( model=model, provider=LlmProviders(_custom_llm_provider), + drop_params=drop_params, ) if provider_config is not None: await base_llm_http_handler.async_realtime( @@ -585,7 +587,6 @@ async def _arealtime( aws_bedrock_runtime_endpoint: Final = kwargs.get("aws_bedrock_runtime_endpoint") aws_external_id: Final = kwargs.get("aws_external_id") aws_session_tags: Final = kwargs.get("aws_session_tags") - drop_params: Final = litellm.drop_params is True or normalize_drop_params(kwargs.get("drop_params")) is True await bedrock_realtime.async_realtime( model=model, @@ -671,6 +672,7 @@ async def _arealtime( resolve_access_token=resolve_vertex_access_token, project=resolved_project, location=vertex_location, + drop_params=drop_params, ) await base_llm_http_handler.async_realtime( diff --git a/litellm/utils.py b/litellm/utils.py index d2ca9faef41..27669d8fd4a 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -9975,6 +9975,7 @@ class ProviderConfigManager: def get_provider_realtime_config( model: str, provider: LlmProviders, + drop_params: bool = False, ) -> BaseRealtimeConfig | None: if LlmProviders.GEMINI == provider: from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig @@ -9983,7 +9984,7 @@ class ProviderConfigManager: if LlmProviders.META == provider: from litellm.llms.meta.realtime.transformation import MetaRealtimeConfig - return MetaRealtimeConfig() + return MetaRealtimeConfig(drop_params=drop_params) return None @staticmethod diff --git a/tests/integration/_support/upstream.py b/tests/integration/_support/upstream.py index d2b67c5438e..186bc345fab 100644 --- a/tests/integration/_support/upstream.py +++ b/tests/integration/_support/upstream.py @@ -9,9 +9,10 @@ import struct import uuid import zlib from collections import deque -from collections.abc import AsyncIterator, Mapping +from collections.abc import AsyncIterator, Callable, Iterator, Mapping from dataclasses import dataclass, field from queue import SimpleQueue +from types import MappingProxyType from typing import Final, cast import httpx @@ -42,7 +43,7 @@ from starlette.datastructures import UploadFile as StarletteUploadFile from starlette.requests import Request from starlette.responses import JSONResponse, Response, StreamingResponse from starlette.routing import Route, WebSocketRoute -from starlette.websockets import WebSocket, WebSocketDisconnect +from starlette.websockets import WebSocket, WebSocketDisconnect, WebSocketState JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) INTERNAL_FIELDS: Final = frozenset( @@ -81,6 +82,25 @@ class Observation: api_key: str = "" +@dataclass(frozen=True, slots=True) +class _RealtimeConnection: + websocket: WebSocket + response: RealtimeResponse + session: dict[str, JsonValue] + scenario_id: str + authorization: str + api_key: str + events: Iterator[dict[str, JsonValue]] + typed_turn_detection: deque[JsonValue] + + def auto_answers(self) -> bool: + return ( + self.response.dialect == "xai" + and len(self.typed_turn_detection) > 0 + and _is_typed_server_vad(self.typed_turn_detection[-1]) + ) + + class InteractionState(BaseModel): """What the scripted Interactions API answers for one interaction id until a DELETE drops it.""" @@ -92,6 +112,12 @@ class InteractionState(BaseModel): delay_seconds: float = 0 +class _ReceivedFrame(BaseModel): + model_config = ConfigDict(frozen=True) + type: str + text: str | None = None + + class _ScenarioRegistration(BaseModel): scenario_id: str response: StoredResponse @@ -415,23 +441,47 @@ class Provider: ) -> None: await websocket.accept() session: Final = _realtime_session(response, scenario_id, websocket.query_params.get("model", "")) + frames: Final = _received_texts(websocket) + held: Final = await anext(frames, None) if response.created_after_first_frame else None for _ in range(response.created_repeats): await websocket.send_json({"type": response.created_event, "event_id": _event_id(), "session": session}) - event_index: Final = iter(response.events) - async for message in websocket.iter_json(): - payload: Final = JSON_OBJECT.validate_python(message) - self.observations.put(Observation(websocket.url.path, authorization, payload, "WEBSOCKET_FRAME", api_key)) - if payload.get("type") == "session.update": - await websocket.send_json( - _realtime_update_reply(payload.get("session"), session, response.session_type) + connection: Final = _RealtimeConnection( + websocket, response, session, scenario_id, authorization, api_key, iter(response.events), deque(maxlen=1) + ) + if held is not None: + await self._realtime_frame(connection, held) + async for message in frames: + await self._realtime_frame(connection, message) + + async def _realtime_frame(self, connection: _RealtimeConnection, message: str) -> None: + websocket: Final = connection.websocket + response: Final = connection.response + payload: Final = JSON_OBJECT.validate_json(message) + self.observations.put( + Observation(websocket.url.path, connection.authorization, payload, "WEBSOCKET_FRAME", connection.api_key) + ) + frame_type: Final = payload.get("type") + if frame_type == "session.update": + reply: Final = _realtime_update_reply( + payload.get("session"), connection.session, response.session_type, response.dialect + ) + if reply["type"] == "session.updated": + connection.typed_turn_detection.extend( + filter(_declares_a_type, _declared_turn_detection(payload.get("session"))) ) - continue - if payload.get("type") not in _REALTIME_TRIGGERS: - continue - event: Final = next(event_index, None) - if event is None: - continue - await websocket.send_json(_rendered_realtime_event(event, scenario_id)) + await _reply_while_open(websocket, reply) + return + if frame_type == "response.cancel": + await _reply_while_open(websocket, _cancel_not_active(payload.get("event_id"), response.dialect)) + return + if frame_type not in _REALTIME_TRIGGERS: + return + event: Final = next(connection.events, None) + if event is None: + return + await _reply_while_open(websocket, _rendered_realtime_event(event, connection.scenario_id)) + if frame_type == "input_audio_buffer.commit" and connection.auto_answers(): + await _reply_while_open(websocket, _rendered_realtime_event(dict(_AUTO_ANSWER), connection.scenario_id)) async def muse_realtime(self, websocket: WebSocket) -> None: await websocket.accept() @@ -454,8 +504,23 @@ class Provider: await websocket.send_json(_rendered_realtime_event(pending.popleft(), scenario_id)) async def gemini_live(self, websocket: WebSocket) -> None: - authorization: Final = websocket.headers.get("authorization", "") - scenario_id: Final = websocket.headers.get("x-goog-user-project", "") + await self._gemini_live_session( + websocket, + websocket.headers.get("x-goog-user-project", ""), + websocket.headers.get("authorization", ""), + _accepts_every_setup, + ) + + async def gemini_ai_studio_live(self, websocket: WebSocket) -> None: + await self._gemini_live_session(websocket, websocket.query_params.get("key", ""), "", _ai_studio_setup_refusal) + + async def _gemini_live_session( + self, + websocket: WebSocket, + scenario_id: str, + authorization: str, + setup_refusal: Callable[[Mapping[str, JsonValue]], str | None], + ) -> None: self.observations.put( Observation( websocket.url.path, @@ -477,6 +542,9 @@ class Provider: Observation(websocket.url.path, authorization, payload, "WEBSOCKET_FRAME", scenario_id) ) if "setup" in payload: + if (refusal := setup_refusal(payload)) is not None: + await websocket.close(code=_INVALID_ARGUMENT_CLOSE, reason=refusal[:_CLOSE_REASON_BYTES]) + return await websocket.send_json({"setupComplete": {}}) continue if not _GEMINI_LIVE_TRIGGERS.intersection(payload): @@ -587,6 +655,7 @@ class Provider: WebSocketRoute("/openai/realtime", self.realtime), WebSocketRoute("/v1/asr/realtime", self.muse_realtime), WebSocketRoute(GEMINI_LIVE_PATH, self.gemini_live), + WebSocketRoute(GEMINI_AI_STUDIO_LIVE_PATH, self.gemini_ai_studio_live), ] ) @@ -612,6 +681,19 @@ def _interaction_body(interaction_id: str, state: InteractionState) -> dict[str, _REALTIME_TRIGGERS: Final = frozenset({"response.create", "input_audio_buffer.commit"}) _GEMINI_LIVE_TRIGGERS: Final = frozenset({"clientContent", "realtimeInput"}) GEMINI_LIVE_PATH: Final = "/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent" +GEMINI_AI_STUDIO_LIVE_PATH: Final = "/ws/google.ai.generativelanguage.v1beta.GenerativeService.BidiGenerateContent" +_INVALID_ARGUMENT_CLOSE: Final = 1007 +_CLOSE_REASON_BYTES: Final = 123 +_AI_STUDIO_TUNING_FIELDS: Final = MappingProxyType( + {"prefixPaddingMs": "prefix_padding_ms", "silenceDurationMs": "silence_duration_ms"} +) +_AUTO_ANSWER: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + "type": "response.done", + "event_id": "evt_$UNIQUE_ID", + "response": {"id": "resp_auto_$UNIQUE_ID", "object": "realtime.response", "status": "completed", "output": []}, + } +) _TRANSCRIPTION_UPDATE_REFUSED: Final = "Passing a realtime session update to a transcription session is not allowed." _REALTIME_UPDATE_REFUSED: Final = "Passing a transcription session update to a realtime session is not allowed." _NESTED_TURN_DETECTION_TYPE: Final = "session.audio.input.turn_detection.type" @@ -651,11 +733,11 @@ def _missing_turn_detection_type(session: Mapping[str, JsonValue]) -> str | None def _realtime_update_reply( - session: JsonValue, created: Mapping[str, JsonValue], session_type: str | None + session: JsonValue, created: Mapping[str, JsonValue], session_type: str | None, dialect: str ) -> dict[str, JsonValue]: if not isinstance(session, dict): return _realtime_error("invalid_value", "Invalid value for 'session': expected an object.", "session") - missing: Final = _missing_turn_detection_type(session) + missing: Final = None if dialect == "xai" else _missing_turn_detection_type(session) if missing is not None: return _realtime_error("missing_required_parameter", f"Missing required parameter: '{missing}'.", missing) declared: Final = session.get("type") @@ -666,6 +748,82 @@ def _realtime_update_reply( return {"type": "session.updated", "event_id": _event_id(), "session": {**created, **session}} +async def _received_texts(websocket: WebSocket) -> AsyncIterator[str]: + while (frame := _ReceivedFrame.model_validate(await websocket.receive())).type == "websocket.receive": + if frame.text is not None: + yield frame.text + + +async def _reply_while_open(websocket: WebSocket, frame: Mapping[str, JsonValue]) -> None: + if websocket.application_state is not WebSocketState.CONNECTED: + return + try: + await websocket.send_json(frame) + except WebSocketDisconnect: + return + + +def _cancel_not_active(cancel_event_id: JsonValue, dialect: str) -> dict[str, JsonValue]: + message: Final = "Cancellation failed: no active response found" + if dialect == "xai": + own_id: Final = _event_id() + return { + "type": "error", + "event_id": own_id, + "error": { + "message": message, + "type": "invalid_request_error", + "code": "invalid_request_error", + "event_id": own_id, + }, + } + return { + "type": "error", + "event_id": _event_id(), + "error": { + "type": "invalid_request_error", + "code": "response_cancel_not_active", + "message": message, + "param": None, + "event_id": cancel_event_id, + }, + } + + +def _declared_turn_detection(session: JsonValue) -> tuple[JsonValue, ...]: + if not isinstance(session, dict): + return () + audio: Final = session.get("audio") + audio_input: Final = audio.get("input") if isinstance(audio, dict) else None + if isinstance(audio_input, dict) and "turn_detection" in audio_input: + return (audio_input["turn_detection"],) + return (session["turn_detection"],) if "turn_detection" in session else () + + +def _declares_a_type(turn_detection: JsonValue) -> bool: + return isinstance(turn_detection, dict) and "type" in turn_detection + + +def _is_typed_server_vad(turn_detection: JsonValue) -> bool: + return isinstance(turn_detection, dict) and turn_detection.get("type") == "server_vad" + + +def _accepts_every_setup(_payload: Mapping[str, JsonValue]) -> str | None: + return None + + +def _ai_studio_setup_refusal(payload: Mapping[str, JsonValue]) -> str | None: + setup: Final = payload.get("setup") + config: Final = setup.get("realtimeInputConfig") if isinstance(setup, dict) else None + detection: Final = config.get("automaticActivityDetection") if isinstance(config, dict) else None + if not isinstance(detection, dict) or detection.get("disabled") is not True: + return None + tuning: Final = tuple(name for key, name in _AI_STUDIO_TUNING_FIELDS.items() if key in detection) + if not tuning: + return None + return f"automatic_activity_detection.disabled is true, but the following fields were also set: {', '.join(tuning)}" + + def _rendered_realtime_event(event: Mapping[str, JsonValue], scenario_id: str) -> dict[str, JsonValue]: return JSON_OBJECT.validate_json( json.dumps(event, separators=(",", ":")) diff --git a/tests/integration/cost_calculation/cost_tracking_case.py b/tests/integration/cost_calculation/cost_tracking_case.py index 038c8d6b9e3..208d95e6a14 100644 --- a/tests/integration/cost_calculation/cost_tracking_case.py +++ b/tests/integration/cost_calculation/cost_tracking_case.py @@ -235,6 +235,8 @@ class RealtimeResponse(BaseModel): session_type: str | None = None created_event: Literal["session.created", "transcription_session.created"] = "session.created" created_repeats: int = 1 + created_after_first_frame: bool = False + dialect: Literal["openai", "xai"] = "openai" StoredResponse: TypeAlias = Annotated[ diff --git a/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py b/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py index f12254b8d1d..ba2261e5a1f 100644 --- a/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py +++ b/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py @@ -8,10 +8,10 @@ while a client frame alone never flips a voice session into one. Every row runs which answers the injected and rewritten updates the way the vendors do. A blocked transcript on a transcription session reaches the client as a ``guardrail_violation`` error and -nothing else: the ``response.cancel``, the voiced block prompt and the ``response.create`` a voice session gets -never go to a backend that cannot speak, while ``on_violation: end_session`` and ``end_session_after_n_fails`` -still close the session. A guardrail that raises anything but a block closes the session the way it does on a -voice session. +nothing else: the ``response.cancel``, the blocked item's deletion, the voiced block prompt and the +``response.create`` a voice session gets never go to a backend that cannot speak, while +``on_violation: end_session`` and ``end_session_after_n_fails`` still close the session. A guardrail that raises +anything but a block closes the session the way it does on a voice session. """ from __future__ import annotations @@ -51,6 +51,7 @@ from tests.integration._support.database import read_rows from tests.integration._support.process import ( OwnedProxy, UpstreamSlot, + graceful_stop_seconds, owned_proxy_process, owned_upstream, stop_root_process, @@ -60,7 +61,7 @@ from tests.integration.cost_calculation.cost_tracking_case import RealtimeRespon RecordProperty = Callable[[str, object], None] -pytestmark: Final = pytest.mark.timeout(180) +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) WORKERS: Final = 2 BURST: Final = 20 @@ -83,7 +84,7 @@ CREATED_TYPES: Final = frozenset({"session.created", "transcription_session.crea TRANSCRIPT_COMPLETED: Final = "conversation.item.input_audio_transcription.completed" RESPONSE_DONE: Final = "response.done" GUARDRAIL_VIOLATION: Final = ("guardrail_violation", "content_policy_violation") -MISSING_TURN_DETECTION_TYPE: Final = ("invalid_request_error", "missing_required_parameter") +SESSION_TYPE_MISMATCH: Final = ("invalid_request_error", "invalid_parameter") INVALID_SESSION_VALUE: Final = ("invalid_request_error", "invalid_value") FIVE_KB: Final = "x" * 5120 UNAUTHENTICATED_STATUS: Final = 403 @@ -104,10 +105,9 @@ BETA_TRANSCRIPTION_UPDATE: Final[dict[str, JsonValue]] = { "input_audio_transcription": {"model": TRANSCRIBE_MODEL}, } VOICE_UPDATE: Final[dict[str, JsonValue]] = {"type": "realtime", "instructions": "answer briefly"} -VOICE_UPDATE_FORWARDED: Final[dict[str, JsonValue]] = { - "type": "realtime", - "instructions": "answer briefly", - "audio": {"input": {"turn_detection": {"create_response": False}}}, +VOICE_UPDATE_GATED: Final[dict[str, JsonValue]] = { + **VOICE_UPDATE, + "audio": {"input": {"turn_detection": {"type": "server_vad", "create_response": False}}}, } RE_ENABLE_UPDATE: Final[dict[str, JsonValue]] = { "type": "realtime", @@ -123,14 +123,32 @@ CLIENT_DECLARED_TRANSCRIPTION: Final[dict[str, JsonValue]] = { } CLIENT_DECLARED_TRANSCRIPTION_FORWARDED: Final[dict[str, JsonValue]] = { "type": "transcription", - "audio": {"input": {"turn_detection": {"create_response": False}}}, + "audio": {"input": {"turn_detection": {"type": "server_vad", "create_response": False}}}, } -TURN_DETECTION_TYPE_MISSING: Final = "Missing required parameter: 'session.audio.input.turn_detection.type'." +REALTIME_UPDATE_REFUSED: Final = "Passing a transcription session update to a realtime session is not allowed." GA_INJECTED_UPDATE: Final[dict[str, JsonValue]] = { "type": "realtime", "audio": {"input": {"turn_detection": {"type": "server_vad", "create_response": False}}}, } BETA_INJECTED_UPDATE: Final[dict[str, JsonValue]] = {"turn_detection": {"type": "server_vad", "create_response": False}} +XAI_VOICE_MODEL: Final = "grok-voice" +GEMINI_LIVE_MODEL: Final = "gemini-3.8-live" +VERTEX_LIVE_LOCATION: Final = "us-central1" +UNTYPED_GATE: Final[dict[str, JsonValue]] = {"create_response": False} +SEMANTIC_VAD: Final[dict[str, JsonValue]] = {"type": "semantic_vad", "eagerness": "low"} +TUNED_SERVER_VAD: Final[dict[str, JsonValue]] = { + "type": "server_vad", + "prefix_padding_ms": 300, + "silence_duration_ms": 500, +} +GEMINI_TUNED_DETECTION: Final[dict[str, JsonValue]] = { + "disabled": True, + "prefixPaddingMs": 300, + "silenceDurationMs": 500, +} +CLIENT_CANCEL_ID: Final = "client_cancel_1" +CANCEL_NOT_ACTIVE: Final = ("invalid_request_error", "response_cancel_not_active") +AUTO_ANSWER_PREFIX: Final = "resp_auto_" COMMIT: Final[dict[str, JsonValue]] = {"type": "input_audio_buffer.commit"} APPEND: Final[dict[str, JsonValue]] = {"type": "input_audio_buffer.append", "audio": PCM_200MS_24KHZ} PUSH_TO_TALK: Final = "PUSH_TO_TALK" @@ -172,7 +190,13 @@ GUARDRAIL_END_CLOSE: Final = 1000 PROXY_FAILURE_CLOSE: Final = 1011 PROXY_FAILURE_REASON: Final = "proxy failed while relaying the upstream websocket" TOOL_OUTPUT_BLOCKED: Final = json.dumps({"error": "Tool output blocked by content policy"}) -VOICE_BLOCK_FRAMES: Final = ("response.cancel", "conversation.item.create", "response.create") +VOICE_BLOCK_FRAMES: Final = ( + "response.cancel", + "conversation.item.delete", + "conversation.item.create", + "response.create", +) +XAI_VOICE_BLOCK_FRAMES: Final = ("response.cancel", "conversation.item.create", "response.create") CREATED_SECONDS: Final = 60.0 STEP_SECONDS: Final = 15.0 SDK_SECONDS: Final = 45.0 @@ -318,6 +342,30 @@ def _voice_scenario(*transcripts: JsonValue) -> RealtimeResponse: return RealtimeResponse(content_type="application/x-realtime", events=tuple(_voice_turns(transcripts))) +def _xai_voice_scenario(transcript: str) -> RealtimeResponse: + return RealtimeResponse( + content_type="application/x-realtime", events=tuple(_voice_turns((transcript,))), dialect="xai" + ) + + +def _voice_scenario_created_after_the_first_frame(transcript: str) -> RealtimeResponse: + return RealtimeResponse( + content_type="application/x-realtime", events=tuple(_voice_turns((transcript,))), created_after_first_frame=True + ) + + +def _live_scenario() -> RealtimeResponse: + return RealtimeResponse(content_type="application/x-realtime", events=()) + + +def _realtime_update(turn_detection: JsonValue) -> dict[str, JsonValue]: + return {"type": "realtime", "audio": {"input": {"turn_detection": turn_detection}}} + + +def _without_auto_response(turn_detection: dict[str, JsonValue]) -> dict[str, JsonValue]: + return {**turn_detection, "create_response": False} + + def _muse_scenario(transcript: str) -> RealtimeResponse: return RealtimeResponse( content_type="application/x-realtime", @@ -358,8 +406,41 @@ def _with_transcription_model(update: dict[str, JsonValue], model: str) -> dict[ return {**update, "audio": {"input": {**audio_input, "transcription": {"model": model}}}} -def _muse_deployment(scenario: Scenario, scenario_id: str, api_base: str) -> str: - return scenario.model(model=f"meta/{MUSE_MODEL}", api_key=scenario_id, api_base=api_base) +def _muse_deployment(scenario: Scenario, scenario_id: str, api_base: str, *, drop_params: bool) -> str: + return scenario.model(model=f"meta/{MUSE_MODEL}", api_key=scenario_id, api_base=api_base, drop_params=drop_params) + + +def _gemini_studio_deployment(scenario: Scenario, scenario_id: str, upstream_url: str) -> str: + return scenario.model( + model=f"gemini/{GEMINI_LIVE_MODEL}", + api_key=scenario_id, + api_base=upstream_url.rstrip("/"), + model_info={"mode": "realtime"}, + ) + + +def _vertex_credentials(upstream_url: str) -> str: + return json.dumps( + { + "type": "external_account", + "audience": "synthetic-vertex-audience", + "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", + "token_url": f"{upstream_url}/_oauth/token", + "credential_source": {"url": f"{upstream_url}/health"}, + } + ) + + +def _vertex_live_deployment(scenario: Scenario, project: str, upstream_url: str) -> str: + upstream: Final = upstream_url.rstrip("/") + return scenario.model( + model=f"vertex_ai/{GEMINI_LIVE_MODEL}", + vertex_project=project, + vertex_credentials=_vertex_credentials(upstream), + vertex_location=VERTEX_LIVE_LOCATION, + api_base=upstream, + model_info={"mode": "realtime"}, + ) def _named_deployment_id(creator: Gateway, scenario: Scenario, name: str, model: str, scenario_id: str) -> str: @@ -549,6 +630,55 @@ def _transcribe( return asyncio.run(_session(ws_base, path, query, key, frames, until=until, headers=headers)) +async def _received_after_sending( + socket: ClientConnection, frames: tuple[dict[str, JsonValue], ...], until: str +) -> AsyncIterator[dict[str, JsonValue]]: + try: + for frame in frames: + await socket.send(json.dumps(frame)) + async for event in _until(socket, until, STEP_SECONDS): + yield event + except ConnectionClosed as closed: + yield {"type": "closed", "code": _close_code(closed)} + + +def _is_proxy_answer(event: dict[str, JsonValue]) -> bool: + if event.get("type") != RESPONSE_DONE: + return False + return not string_value(object_value(event["response"])["id"]).startswith(AUTO_ANSWER_PREFIX) + + +async def _until_the_proxy_answer( + socket: ClientConnection, frames: tuple[dict[str, JsonValue], ...] +) -> AsyncIterator[dict[str, JsonValue]]: + try: + yield JSON_OBJECT.validate_json(await asyncio.wait_for(socket.recv(), CREATED_SECONDS)) + for frame in frames: + await socket.send(json.dumps(frame)) + deadline: Final = asyncio.get_running_loop().time() + STEP_SECONDS + while (event := await _next_event(socket, deadline)) is not None: + yield event + if _is_proxy_answer(event): + return + yield {"type": "timeout"} + except ConnectionClosed as closed: + yield {"type": "closed", "code": _close_code(closed)} + + +async def _talked_until_the_proxy_answer( + ws_base: str, query: str, key: str, frames: tuple[dict[str, JsonValue], ...] +) -> Session: + headers: Final = {"Authorization": f"Bearer {key}"} + async with websockets.connect(f"{ws_base}/v1/realtime?{query}", additional_headers=headers) as socket: + return Session(tuple([event async for event in _until_the_proxy_answer(socket, frames)]), None) + + +async def _talked_at_connect(ws_base: str, query: str, key: str, frames: tuple[dict[str, JsonValue], ...]) -> Session: + headers: Final = {"Authorization": f"Bearer {key}"} + async with websockets.connect(f"{ws_base}/v1/realtime?{query}", additional_headers=headers) as socket: + return Session(tuple([event async for event in _received_after_sending(socket, frames, RESPONSE_DONE)]), None) + + def _talk( ws_base: str, query: str, @@ -617,10 +747,43 @@ def _session_updates(observed: tuple[dict[str, JsonValue], ...]) -> tuple[JsonVa return tuple(frame.get("session") for frame in _sent(observed) if frame.get("type") == "session.update") +def _assert_voice_updates(observed: tuple[dict[str, JsonValue], ...], *later: dict[str, JsonValue]) -> None: + updates: Final = _session_updates(observed) + assert updates[0] == GA_INJECTED_UPDATE, updates + assert updates[1] == VOICE_UPDATE_GATED, updates + assert updates[2:] == later, updates + + def _query(upgrade: dict[str, JsonValue]) -> JsonValue: return upgrade["query"] +def _declared_turn_detection(update: JsonValue) -> JsonValue: + session: Final = object_value(update) + audio_input: Final = object_value(object_value(session.get("audio", {})).get("input", {})) + return audio_input["turn_detection"] if "turn_detection" in audio_input else session.get("turn_detection") + + +def _upstream_turn_detections(observed: tuple[dict[str, JsonValue], ...]) -> tuple[JsonValue, ...]: + return tuple(_declared_turn_detection(update) for update in _session_updates(observed)) + + +def _response_ids(session: Session) -> tuple[str, ...]: + done: Final = tuple(object_value(event["response"]) for event in session.events if event["type"] == RESPONSE_DONE) + return tuple(string_value(response["id"]) for response in done) + + +def _error_event_ids(session: Session) -> tuple[JsonValue, ...]: + errors: Final = tuple(object_value(event["error"]) for event in session.events if event["type"] == "error") + return tuple(error.get("event_id") for error in errors) + + +def _setup_detection(observed: tuple[dict[str, JsonValue], ...]) -> JsonValue: + setups: Final = tuple(object_value(frame["setup"]) for frame in _sent(observed) if "setup" in frame) + assert len(setups) == 1, _sent(observed) + return object_value(setups[0].get("realtimeInputConfig", {})).get("automaticActivityDetection") + + def _spend_rows(key: str, count: int) -> list[dict[str, JsonValue]]: return eventually( lambda: read_rows(SPEND_SQL, (sha256(key.encode()).hexdigest(),)), @@ -704,6 +867,7 @@ def _write_config(directory: Path, ca_bundle: Path) -> Path: "litellm_settings": { "ssl_verify": str(ca_bundle), "enable_redis_auth_cache": True, + "gemini_live_defer_setup": True, "cache": True, "cache_params": {"type": "redis", "host": "os.environ/REDIS_HOST", "port": "os.environ/REDIS_PORT"}, }, @@ -918,19 +1082,18 @@ def test_voice_session_keeps_the_injected_update_and_the_clean_transcript_flows( assert session.types == ( "session.created", "session.updated", - "error", + "session.updated", TRANSCRIPT_COMPLETED, RESPONSE_DONE, ), session - assert session.errors == (MISSING_TURN_DETECTION_TYPE,), session - assert session.error_messages == (TURN_DETECTION_TYPE_MISSING,), session + assert session.errors == (), session assert _sent_types(observed) == ( "session.update", "session.update", "input_audio_buffer.commit", "response.create", ), _sent(observed) - assert _session_updates(observed) == (GA_INJECTED_UPDATE, VOICE_UPDATE_FORWARDED), _session_updates(observed) + _assert_voice_updates(observed) rows: Final = _spend_rows(key, 1) assert rows[0]["call_type"] == "_arealtime", rows @@ -945,17 +1108,15 @@ def test_voice_session_blocked_transcript_is_refused_through_the_backend(guardra ) observed: Final = _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id) assert session.transcripts == (BLOCKED_TRANSCRIPT,), session - assert session.errors == (MISSING_TURN_DETECTION_TYPE, GUARDRAIL_VIOLATION), session + assert session.errors == (GUARDRAIL_VIOLATION,), session assert session.types[-1] == RESPONSE_DONE, session assert _sent_types(observed) == ( "session.update", "session.update", "input_audio_buffer.commit", - "response.cancel", - "conversation.item.create", - "response.create", + *VOICE_BLOCK_FRAMES, ), _sent(observed) - assert _session_updates(observed) == (GA_INJECTED_UPDATE, VOICE_UPDATE_FORWARDED), _session_updates(observed) + _assert_voice_updates(observed) def test_voice_session_cannot_re_enable_auto_response_with_a_later_update(guardrail_proxy: OwnedProxy) -> None: @@ -967,12 +1128,8 @@ def test_voice_session_cannot_re_enable_auto_response_with_a_later_update(guardr session: Final = _talk(_ws_base(_owned_url(guardrail_proxy)), f"model={model}", key, frames) observed: Final = _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id) assert session.types[-1] == RESPONSE_DONE, session - assert session.errors == (MISSING_TURN_DETECTION_TYPE,), session - assert _session_updates(observed) == ( - GA_INJECTED_UPDATE, - VOICE_UPDATE_FORWARDED, - RE_ENABLE_FORWARDED, - ), _session_updates(observed) + assert session.errors == (), session + _assert_voice_updates(observed, RE_ENABLE_FORWARDED) def test_client_declared_transcription_type_does_not_bypass_the_voice_guardrail(guardrail_proxy: OwnedProxy) -> None: @@ -984,16 +1141,14 @@ def test_client_declared_transcription_type_does_not_bypass_the_voice_guardrail( session: Final = _talk(_ws_base(_owned_url(guardrail_proxy)), f"model={model}", key, frames, seconds=20) observed: Final = _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id) assert session.transcripts == (BLOCKED_TRANSCRIPT,), session - assert session.errors == (MISSING_TURN_DETECTION_TYPE, GUARDRAIL_VIOLATION), session - assert session.error_messages[0] == TURN_DETECTION_TYPE_MISSING, session + assert session.errors == (SESSION_TYPE_MISMATCH, GUARDRAIL_VIOLATION), session + assert session.error_messages[0] == REALTIME_UPDATE_REFUSED, session assert session.types[-1] == RESPONSE_DONE, session assert _sent_types(observed) == ( "session.update", "session.update", "input_audio_buffer.commit", - "response.cancel", - "conversation.item.create", - "response.create", + *VOICE_BLOCK_FRAMES, ), _sent(observed) assert _session_updates(observed) == ( GA_INJECTED_UPDATE, @@ -1014,6 +1169,202 @@ def test_backend_session_created_typed_transcription_skips_the_injection_on_the_ assert [_query(upgrade) for upgrade in _upgrades(observed)] == [[["model", TRANSCRIBE_MODEL]]] +def test_voice_session_typed_turn_detection_keeps_its_type_and_tuning(guardrail_proxy: OwnedProxy) -> None: + with guardrail_proxy.gateway.scenario() as scenario: + handle: Final = _scripted(scenario, _voice_scenario(CLEAN_TRANSCRIPT)) + key: Final = scenario.key() + model: Final = _openai_deployment(scenario, handle.scenario_id, model=VOICE_MODEL) + frames: Final = _voice_frames(_realtime_update(SEMANTIC_VAD)) + session: Final = _talk(_ws_base(_owned_url(guardrail_proxy)), f"model={model}", key, frames) + observed: Final = _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id) + assert session.errors == (), session + assert session.types[-1] == RESPONSE_DONE, session + assert _session_updates(observed) == ( + GA_INJECTED_UPDATE, + _realtime_update(_without_auto_response(SEMANTIC_VAD)), + ), _session_updates(observed) + + +@dataclass(frozen=True, slots=True) +class LaterUpdateCase: + headers: Mapping[str, str] + first: dict[str, JsonValue] + later: dict[str, JsonValue] + forwarded: tuple[dict[str, JsonValue], ...] + + +LATER_UPDATE_CASES: Final = ( + pytest.param( + LaterUpdateCase( + {}, + VOICE_UPDATE, + {"type": "realtime", "instructions": "later"}, + (GA_INJECTED_UPDATE, VOICE_UPDATE_GATED), + ), + id="ga", + ), + pytest.param( + LaterUpdateCase( + BETA_HEADERS, + {"turn_detection": {"type": "server_vad"}}, + {"instructions": "later"}, + (BETA_INJECTED_UPDATE, BETA_INJECTED_UPDATE), + ), + id="beta", + ), +) + + +@pytest.mark.parametrize("case", LATER_UPDATE_CASES) +def test_later_update_without_turn_detection_is_forwarded_verbatim( + guardrail_proxy: OwnedProxy, case: LaterUpdateCase +) -> None: + with guardrail_proxy.gateway.scenario() as scenario: + handle: Final = _scripted(scenario, _voice_scenario(CLEAN_TRANSCRIPT)) + key: Final = scenario.key() + model: Final = _openai_deployment(scenario, handle.scenario_id, model=VOICE_MODEL) + session: Final = _drive( + _ws_base(_owned_url(guardrail_proxy)), + f"model={model}", + key, + (_step(_update_frame(case.first), _update_frame(case.later), COMMIT, until=RESPONSE_DONE),), + headers=case.headers, + ) + observed: Final = _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id) + assert session.errors == (), session + assert session.types[-1] == RESPONSE_DONE, session + assert _session_updates(observed) == (*case.forwarded, case.later), _session_updates(observed) + + +def test_session_update_sent_before_session_created_follows_the_gate(guardrail_proxy: OwnedProxy) -> None: + with guardrail_proxy.gateway.scenario() as scenario: + handle: Final = _scripted(scenario, _voice_scenario_created_after_the_first_frame(CLEAN_TRANSCRIPT)) + key: Final = scenario.key() + model: Final = _openai_deployment(scenario, handle.scenario_id, model=VOICE_MODEL) + frames: Final = _voice_frames(_realtime_update(TUNED_SERVER_VAD)) + session: Final = asyncio.run( + _talked_at_connect(_ws_base(_owned_url(guardrail_proxy)), f"model={model}", key, frames) + ) + observed: Final = _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id) + assert session.types == ( + "session.created", + "session.updated", + "session.updated", + TRANSCRIPT_COMPLETED, + RESPONSE_DONE, + ), session + assert session.errors == (), session + assert _session_updates(observed) == ( + GA_INJECTED_UPDATE, + _realtime_update(_without_auto_response(TUNED_SERVER_VAD)), + ), _session_updates(observed) + + +def test_client_response_cancel_error_still_reaches_the_client(guardrail_proxy: OwnedProxy) -> None: + with guardrail_proxy.gateway.scenario() as scenario: + handle: Final = _scripted(scenario, _voice_scenario(CLEAN_TRANSCRIPT)) + key: Final = scenario.key() + model: Final = _openai_deployment(scenario, handle.scenario_id, model=VOICE_MODEL) + cancel: Final[dict[str, JsonValue]] = {"type": "response.cancel", "event_id": CLIENT_CANCEL_ID} + session: Final = _drive( + _ws_base(_owned_url(guardrail_proxy)), f"model={model}", key, (_step(cancel, until="error"),) + ) + assert session.errors == (CANCEL_NOT_ACTIVE,), session + assert _error_event_ids(session) == (CLIENT_CANCEL_ID,), session + + +def _xai_blocked_session( + guardrail_proxy: OwnedProxy, scenario: Scenario +) -> tuple[Session, tuple[dict[str, JsonValue], ...]]: + handle: Final = _scripted(scenario, _xai_voice_scenario(BLOCKED_TRANSCRIPT)) + key: Final = scenario.key() + model: Final = scenario.model(model=f"xai/{XAI_VOICE_MODEL}", api_key=handle.scenario_id) + frames: Final = _voice_frames(_realtime_update(TUNED_SERVER_VAD)) + session: Final = asyncio.run( + _talked_until_the_proxy_answer(_ws_base(_owned_url(guardrail_proxy)), f"model={model}", key, frames) + ) + return session, _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id) + + +def test_xai_voice_session_gets_untyped_gates_and_no_unguarded_answer(guardrail_proxy: OwnedProxy) -> None: + with guardrail_proxy.gateway.scenario() as scenario: + session, observed = _xai_blocked_session(guardrail_proxy, scenario) + turn_detections: Final = _upstream_turn_detections(observed) + assert turn_detections == (UNTYPED_GATE, UNTYPED_GATE), turn_detections + assert session.transcripts == (BLOCKED_TRANSCRIPT,), session + assert GUARDRAIL_VIOLATION in session.errors, session + assert not any(response.startswith(AUTO_ANSWER_PREFIX) for response in _response_ids(session)), session + + +def test_xai_voice_session_block_sends_xai_no_item_delete(guardrail_proxy: OwnedProxy) -> None: + with guardrail_proxy.gateway.scenario() as scenario: + session, observed = _xai_blocked_session(guardrail_proxy, scenario) + assert session.transcripts == (BLOCKED_TRANSCRIPT,), session + assert _sent_types(observed) == ( + "session.update", + "session.update", + "input_audio_buffer.commit", + *XAI_VOICE_BLOCK_FRAMES, + ), _sent(observed) + + +@pytest.mark.xfail( + strict=True, + raises=AssertionError, + reason=( + "xAI answers the guardrail's response.cancel with an error carrying its own event_id, so the proxy cannot " + "tell it apart from an error for a client event and forwards it, as the merge base does" + ), +) +def test_xai_guardrail_cancel_error_does_not_reach_the_client(guardrail_proxy: OwnedProxy) -> None: + with guardrail_proxy.gateway.scenario() as scenario: + session, _ = _xai_blocked_session(guardrail_proxy, scenario) + assert session.errors == (GUARDRAIL_VIOLATION,), session + + +TUNED_VOICE_UPDATES: Final = ( + pytest.param(_realtime_update(TUNED_SERVER_VAD), id="nested"), + pytest.param({"turn_detection": TUNED_SERVER_VAD}, id="flat"), +) + + +@pytest.mark.parametrize("update", TUNED_VOICE_UPDATES) +def test_gemini_ai_studio_voice_session_with_vad_tuning_connects_under_the_guardrail( + guardrail_proxy: OwnedProxy, update: dict[str, JsonValue] +) -> None: + with guardrail_proxy.gateway.scenario() as scenario: + handle: Final = _scripted(scenario, _live_scenario()) + key: Final = scenario.key() + upstream_url: Final = guardrail_proxy.gateway.upstream_url + model: Final = _gemini_studio_deployment(scenario, handle.scenario_id, upstream_url) + session: Final = _drive( + _ws_base(_owned_url(guardrail_proxy)), + f"model={model}", + key, + (_step(_update_frame(update), until="session.updated"),), + ) + observed: Final = _observed(upstream_url, handle.scenario_id) + assert session.types == ("session.created", "session.updated"), session + assert _setup_detection(observed) == {"disabled": True}, _sent(observed) + + +def test_vertex_live_voice_session_keeps_vad_tuning_under_the_guardrail(guardrail_proxy: OwnedProxy) -> None: + with guardrail_proxy.gateway.scenario() as scenario: + handle: Final = _scripted(scenario, _live_scenario()) + key: Final = scenario.key() + upstream_url: Final = guardrail_proxy.gateway.upstream_url + model: Final = _vertex_live_deployment(scenario, handle.scenario_id, upstream_url) + session: Final = _drive( + _ws_base(_owned_url(guardrail_proxy)), + f"model={model}", + key, + (_step(_update_frame(_realtime_update(TUNED_SERVER_VAD)), until="session.updated"),), + ) + observed: Final = _observed(upstream_url, handle.scenario_id) + assert session.types == ("session.created", "session.updated"), session + assert _setup_detection(observed) == GEMINI_TUNED_DETECTION, _sent(observed) + + MUSE_PUSH_TO_TALK_FRAMES: Final[tuple[dict[str, JsonValue], ...]] = ( {"binary_bytes": MUSE_PACKET_BYTES}, {"binary_bytes": MUSE_PACKET_BYTES}, @@ -1041,10 +1392,11 @@ def _muse_session( turn_detection: JsonValue, *, until: str | None = None, + drop_params: bool = False, ) -> tuple[Session, tuple[dict[str, JsonValue], ...]]: handle: Final = _scripted(scenario, _muse_scenario(transcript), control_url=tls_upstream) key: Final = scenario.key() - model: Final = _muse_deployment(scenario, handle.scenario_id, tls_upstream) + model: Final = _muse_deployment(scenario, handle.scenario_id, tls_upstream, drop_params=drop_params) verdict: Final = TRANSCRIPT_COMPLETED if BLOCKED_WORD not in transcript else "error" session: Final = _talk( _ws_base(_owned_url(guardrail_proxy)), @@ -1120,6 +1472,40 @@ def test_muse_server_vad_transcription_session_keeps_endpointing( ), _muse_audio_frames(observed) +def test_muse_semantic_vad_under_drop_params_falls_back_to_endpointing( + guardrail_proxy: OwnedProxy, tls_upstream: str +) -> None: + with guardrail_proxy.gateway.scenario() as scenario: + session, observed = _muse_session( + guardrail_proxy, + scenario, + tls_upstream, + CLEAN_TRANSCRIPT, + f"&{TRANSCRIPTION_QUERY}", + {"type": "semantic_vad"}, + drop_params=True, + ) + assert session.transcripts == (CLEAN_TRANSCRIPT,), session + assert session.errors == (), session + assert _muse_handshake(observed)["mode"] == ENDPOINTING, _muse_handshake(observed) + + +def test_muse_semantic_vad_without_drop_params_is_refused(guardrail_proxy: OwnedProxy, tls_upstream: str) -> None: + with guardrail_proxy.gateway.scenario() as scenario: + session, observed = _muse_session( + guardrail_proxy, + scenario, + tls_upstream, + CLEAN_TRANSCRIPT, + "", + {"type": "semantic_vad"}, + until=TRANSCRIPT_COMPLETED, + ) + assert session.types[-1] == "closed", session + assert session.transcripts == (), session + assert _upgrades(observed) == (), observed + + @pytest.mark.xfail( strict=True, reason=( @@ -1229,8 +1615,8 @@ def test_opt_in_guardrail_gates_a_voice_session_on_an_opted_out_key(guardrail_pr ) observed: Final = _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id) assert session.transcripts == (BLOCKED_TRANSCRIPT,), session - assert session.errors == (MISSING_TURN_DETECTION_TYPE, GUARDRAIL_VIOLATION), session - assert _session_updates(observed) == (GA_INJECTED_UPDATE, VOICE_UPDATE_FORWARDED), _session_updates(observed) + assert session.errors == (GUARDRAIL_VIOLATION,), session + _assert_voice_updates(observed) def test_opt_in_guardrail_leaves_a_transcription_session_alone_on_an_opted_out_key(guardrail_proxy: OwnedProxy) -> None: @@ -1717,12 +2103,11 @@ def test_voice_session_guardrail_raising_runtime_error_closes_with_the_same_prox assert session.types == ( "session.created", "session.updated", - "error", + "session.updated", TRANSCRIPT_COMPLETED, "error", "closed", ), session - assert session.errors[0] == MISSING_TURN_DETECTION_TYPE, session _assert_closed_by_a_proxy_failure(session) assert _sent_types(observed) == ("session.update", "session.update", "input_audio_buffer.commit"), _sent( observed @@ -1740,8 +2125,8 @@ def test_voice_session_guardrail_raising_value_error_is_voiced_through_the_backe (_step(VOICE_UPDATE_FRAME, COMMIT, until=RESPONSE_DONE),), ) assert session.transcripts == (blocked,), session - assert session.errors == (MISSING_TURN_DETECTION_TYPE, GUARDRAIL_VIOLATION), session - assert session.error_messages[1] == f"{VALUE_ERROR_WORD} is not allowed", session + assert session.errors == (GUARDRAIL_VIOLATION,), session + assert session.error_messages == (f"{VALUE_ERROR_WORD} is not allowed",), session assert session.types[-1] == RESPONSE_DONE, session assert _sent_types(observed) == ( "session.update", @@ -1764,7 +2149,7 @@ def test_voice_session_second_violation_under_end_session_after_n_fails_closes_a (_step(VOICE_UPDATE_FRAME, COMMIT, until=RESPONSE_DONE), _step(COMMIT, until="closed")), ) assert session.transcripts == (blocked, blocked), session - assert session.errors == (MISSING_TURN_DETECTION_TYPE, GUARDRAIL_VIOLATION, GUARDRAIL_VIOLATION), session + assert session.errors == (GUARDRAIL_VIOLATION, GUARDRAIL_VIOLATION), session assert session.types[-1] == "closed", session assert session.close_code == GUARDRAIL_END_CLOSE, session assert _sent_types(observed) == ( @@ -1948,7 +2333,7 @@ def test_non_string_transcript_field_closes_a_voice_session_with_the_same_proxy_ assert session.types == ( "session.created", "session.updated", - "error", + "session.updated", TRANSCRIPT_COMPLETED, "error", "closed", diff --git a/tests/integration/providers/test_realtime_unguarded_voice_turns.py b/tests/integration/providers/test_realtime_unguarded_voice_turns.py index 4dd7fd9aeef..c7489685273 100644 --- a/tests/integration/providers/test_realtime_unguarded_voice_turns.py +++ b/tests/integration/providers/test_realtime_unguarded_voice_turns.py @@ -114,12 +114,10 @@ TRANSCRIPT_FIELDS: Final = ( pytest.param({"transcript": {"text": "a"}}, id="object"), ) VOICE_UPDATE: Final[dict[str, JsonValue]] = {"type": "realtime", "instructions": "answer briefly"} -VOICE_UPDATE_FORWARDED: Final[dict[str, JsonValue]] = { - "type": "realtime", - "instructions": "answer briefly", - "audio": {"input": {"turn_detection": {"create_response": False}}}, +VOICE_UPDATE_GATED: Final[dict[str, JsonValue]] = { + **VOICE_UPDATE, + "audio": {"input": {"turn_detection": {"type": "server_vad", "create_response": False}}}, } -TURN_DETECTION_TYPE_MISSING: Final = "Missing required parameter: 'session.audio.input.turn_detection.type'." INJECTED_UPDATE: Final[dict[str, JsonValue]] = { "type": "realtime", "audio": {"input": {"turn_detection": {"type": "server_vad", "create_response": False}}}, @@ -982,15 +980,15 @@ def test_list_mode_transcript_guardrail_still_turns_auto_response_off_and_answer assert session.types == ( "session.created", "session.updated", - "error", + "session.updated", TRANSCRIPT_COMPLETED, RESPONSE_DONE, ), session assert session.transcript_fields == ({"transcript": TRANSCRIPT},), session - assert session.error_messages == (TURN_DETECTION_TYPE_MISSING,), session + assert session.error_messages == (), session assert _sent(observed) == ( _update_frame(INJECTED_UPDATE), - _update_frame(VOICE_UPDATE_FORWARDED), + _update_frame(VOICE_UPDATE_GATED), COMMIT, PROXY_RESPONSE_CREATE, ), _sent(observed) diff --git a/tests/unit/integration_support/test_upstream.py b/tests/unit/integration_support/test_upstream.py index 0cbcf1b7ee1..9e114e3b0b7 100644 --- a/tests/unit/integration_support/test_upstream.py +++ b/tests/unit/integration_support/test_upstream.py @@ -21,7 +21,11 @@ async def test_realtime_disconnect_after_session_creation_keeps_unexpected_error response: Final = RealtimeResponse(content_type="application/x-realtime", events=({"type": "response.done"},)) provider.scenario_store.put("owned-scenario", response) incoming: Final[deque[Message]] = deque( - ({"type": "websocket.connect"}, {"type": "websocket.receive", "text": '{"type":"response.create"}'}) + ( + {"type": "websocket.connect"}, + {"type": "websocket.receive", "text": '{"type":"response.create"}'}, + {"type": "websocket.disconnect", "code": 1006}, + ) ) outgoing: Final[deque[Message]] = deque() diff --git a/tests/unit/litellm_core_utils/test_realtime_streaming.py b/tests/unit/litellm_core_utils/test_realtime_streaming.py index 8de5e27d99c..9b2c42450ae 100644 --- a/tests/unit/litellm_core_utils/test_realtime_streaming.py +++ b/tests/unit/litellm_core_utils/test_realtime_streaming.py @@ -2351,8 +2351,7 @@ async def test_guardrail_turn_detection_injection_tolerates_non_dict_value( injected_turn_detection = session_obj.get("turn_detection") or session_obj.get("audio", {}).get("input", {}).get( "turn_detection" ) - assert isinstance(injected_turn_detection, dict) - assert injected_turn_detection["create_response"] is False + assert injected_turn_detection == {"type": "server_vad", "create_response": False} assert streaming._guardrail_turn_detection_update_sent is True @@ -2417,8 +2416,7 @@ async def test_subsequent_session_update_cannot_reenable_vad_when_guardrails_act forwarded_turn_detection = session_obj.get("turn_detection") or session_obj.get("audio", {}).get("input", {}).get( "turn_detection" ) - assert isinstance(forwarded_turn_detection, dict) - assert forwarded_turn_detection["create_response"] is False + assert forwarded_turn_detection == {"type": "server_vad", "create_response": False} @pytest.mark.asyncio @@ -3927,3 +3925,317 @@ async def test_transcription_session_guardrail_block_only_reports_violation( sent_to_backend: Final = [call.args[0] for call in backend_ws.send.await_args_list] assert sent_to_backend == [], sent_to_backend assert backend_ws.close.await_count == (1 if expect_session_closed else 0) + + +_GA_GUARDRAIL_GATE: Final = { + "type": "realtime", + "audio": {"input": {"turn_detection": {"type": "server_vad", "create_response": False}}}, +} + + +async def _forwarded_ga_session_updates( + monkeypatch: pytest.MonkeyPatch, + *client_sessions: dict[str, object], + event_normalizer: XAIRealtimeNormalizer | None = None, +) -> tuple[dict[str, object], ...]: + """Run a GA client's session.updates through a transcript-guardrailed proxy and return the forwarded sessions.""" + monkeypatch.setattr(litellm, "callbacks", [_transcription_guardrail()]) + client_ws: Final = AsyncMock() + client_ws.scope = {"headers": []} + client_ws.receive_text = AsyncMock( + side_effect=[ + *(json.dumps({"type": "session.update", "session": session}) for session in client_sessions), + ConnectionClosed(None, None), + ] + ) + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + logging_obj: Final = MagicMock() + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + await RealTimeStreaming(client_ws, backend_ws, logging_obj, event_normalizer=event_normalizer).client_ack_messages() + return tuple(json.loads(call.args[0])["session"] for call in backend_ws.send.await_args_list) + + +@pytest.mark.asyncio +async def test_first_ga_session_update_keeps_client_turn_detection_when_guardrail_disables_auto_response( + monkeypatch: pytest.MonkeyPatch, +): + """A transcript guardrail only flips create_response: the client's nested GA + turn_detection, its type included, must reach the backend intact, or OpenAI + rejects the whole update and the client's session config is lost.""" + forwarded: Final = await _forwarded_ga_session_updates( + monkeypatch, + { + "type": "realtime", + "instructions": "hi", + "audio": { + "input": { + "turn_detection": {"type": "server_vad", "silence_duration_ms": 700, "create_response": True}, + "transcription": {"model": "gpt-4o-transcribe"}, + } + }, + }, + ) + + gate, client_update = forwarded + assert gate == _GA_GUARDRAIL_GATE + assert "turn_detection" not in client_update + assert client_update["instructions"] == "hi" + assert client_update["audio"]["input"] == { + "turn_detection": {"type": "server_vad", "silence_duration_ms": 700, "create_response": False}, + "transcription": {"model": "gpt-4o-transcribe"}, + } + + +@pytest.mark.asyncio +async def test_first_ga_session_update_without_turn_detection_gets_a_typed_gate(monkeypatch: pytest.MonkeyPatch): + forwarded: Final = await _forwarded_ga_session_updates(monkeypatch, {"type": "realtime", "instructions": "hi"}) + + gate, client_update = forwarded + assert gate == _GA_GUARDRAIL_GATE + assert "turn_detection" not in client_update + assert client_update["audio"]["input"]["turn_detection"] == {"type": "server_vad", "create_response": False} + + +@pytest.mark.asyncio +async def test_later_ga_session_update_without_turn_detection_is_forwarded_without_one( + monkeypatch: pytest.MonkeyPatch, +): + """OpenAI keeps the earlier create_response=false for a field left unset, while a + bare create_response carries no type and gets the whole update rejected.""" + forwarded: Final = await _forwarded_ga_session_updates( + monkeypatch, + {"type": "realtime", "audio": {"input": {"turn_detection": {"type": "server_vad", "create_response": True}}}}, + {"type": "realtime", "instructions": "later"}, + ) + + gate, first, later = forwarded + assert gate == _GA_GUARDRAIL_GATE + assert first["audio"]["input"]["turn_detection"] == {"type": "server_vad", "create_response": False} + assert later == {"type": "realtime", "instructions": "later"} + + +_XAI_TYPED_VAD_SESSION: Final = { + "type": "realtime", + "audio": {"input": {"turn_detection": {"type": "server_vad", "silence_duration_ms": 700}}}, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "client_sessions", + [ + pytest.param(({"type": "realtime", "instructions": "hi"},), id="no-turn-detection"), + pytest.param(({"type": "realtime", "audio": {"input": {"turn_detection": None}}},), id="push-to-talk"), + pytest.param((_XAI_TYPED_VAD_SESSION,), id="typed-server-vad"), + pytest.param( + ({"type": "realtime", "instructions": "hi"}, _XAI_TYPED_VAD_SESSION), id="later-typed-server-vad" + ), + ], +) +async def test_xai_session_update_gate_leaves_out_the_vad_type( + monkeypatch: pytest.MonkeyPatch, client_sessions: tuple[dict[str, object], ...] +): + """xAI answers on its own under a typed turn_detection even with create_response off, + so every turn_detection the proxy forwards to xAI under the guardrail must carry no type.""" + forwarded: Final = await _forwarded_ga_session_updates( + monkeypatch, *client_sessions, event_normalizer=XAIRealtimeNormalizer() + ) + + assert [session["audio"]["input"]["turn_detection"] for session in forwarded] == [ + {"create_response": False} + ] * (len(client_sessions) + 1) + + +@pytest.mark.asyncio +async def test_xai_session_created_gate_leaves_out_the_vad_type(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "callbacks", [_transcription_guardrail()]) + client_ws: Final = MagicMock() + client_ws.send_text = AsyncMock() + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + backend_ws.recv = AsyncMock( + side_effect=[ + json.dumps({"type": "session.created", "session": {"turn_detection": {"type": None}}}).encode(), + ConnectionClosed(None, None), + ] + ) + logging_obj: Final = _mock_logging() + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + streaming: Final = RealTimeStreaming(client_ws, backend_ws, logging_obj, event_normalizer=XAIRealtimeNormalizer()) + + await streaming.backend_to_client_send_messages() + + sent: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] + assert sent == [ + { + "type": "session.update", + "session": {"type": "realtime", "audio": {"input": {"turn_detection": {"create_response": False}}}}, + } + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "client_turn_detection", + [ + pytest.param({"type": "semantic_vad", "eagerness": "low"}, id="semantic-vad"), + pytest.param({"type": "server_vad", "silence_duration_ms": 700}, id="tuned-server-vad"), + ], +) +async def test_guardrail_gate_goes_ahead_of_a_session_update_sent_before_session_created( + monkeypatch: pytest.MonkeyPatch, client_turn_detection: dict[str, object] +): + """OpenAI applies events in order, so the gate must precede a client update sent on connect: + the client's turn_detection then wins, and when OpenAI refuses that update the gate still holds.""" + monkeypatch.setattr(litellm, "callbacks", [_transcription_guardrail()]) + client_ws: Final = AsyncMock() + client_ws.scope = {"headers": []} + client_ws.receive_text = AsyncMock( + side_effect=[ + json.dumps( + { + "type": "session.update", + "session": {"type": "realtime", "audio": {"input": {"turn_detection": client_turn_detection}}}, + } + ), + ConnectionClosed(None, None), + ] + ) + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + backend_ws.recv = AsyncMock( + side_effect=[json.dumps({"type": "session.created", "session": {}}).encode(), ConnectionClosed(None, None)] + ) + logging_obj: Final = _mock_logging() + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + streaming: Final = RealTimeStreaming(client_ws, backend_ws, logging_obj) + + await streaming.client_ack_messages() + await streaming.backend_to_client_send_messages() + + sent: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] + assert [event["session"]["audio"]["input"]["turn_detection"] for event in sent] == [ + {"type": "server_vad", "create_response": False}, + {**client_turn_detection, "create_response": False}, + ] + + +@pytest.mark.asyncio +async def test_client_never_sees_the_error_caused_by_the_guardrail_cancel(monkeypatch: pytest.MonkeyPatch): + """OpenAI answers the guardrail's response.cancel with response_cancel_not_active when nothing + is answering yet, an error about an event the client never sent.""" + + class _BlockingGuardrail(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + raise ValueError("blocked") + + monkeypatch.setattr( + litellm, + "callbacks", + [ + _BlockingGuardrail( + guardrail_name="blocking", event_hook=GuardrailEventHooks.realtime_input_transcription, default_on=True + ) + ], + ) + client_ws: Final = MagicMock() + client_ws.send_text = AsyncMock() + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + + def cancel_not_active(event_id: str) -> bytes: + return json.dumps( + {"type": "error", "error": {"code": "response_cancel_not_active", "event_id": event_id}} + ).encode() + + def proxy_cancel_event_id() -> str: + sent = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] + return next(event["event_id"] for event in sent if event["type"] == "response.cancel") + + backend_frames: Final = iter( + ( + lambda: _make_transcript_event("blocked words"), + lambda: cancel_not_active(proxy_cancel_event_id()), + lambda: cancel_not_active("evt_from_client"), + lambda: cancel_not_active("litellm_guardrail_from_client"), + ) + ) + + async def recv(**_kwargs: object) -> bytes: + frame = next(backend_frames, None) + if frame is None: + raise ConnectionClosed(None, None) + return frame() + + backend_ws.recv = recv + logging_obj: Final = _mock_logging() + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + + await RealTimeStreaming(client_ws, backend_ws, logging_obj).backend_to_client_send_messages() + + client_errors: Final = [ + event["error"] for call in client_ws.send_text.await_args_list if (event := json.loads(call.args[0]))["type"] == "error" + ] + assert [(error["code"], error.get("event_id")) for error in client_errors] == [ + ("content_policy_violation", None), + ("response_cancel_not_active", "evt_from_client"), + ("response_cancel_not_active", "litellm_guardrail_from_client"), + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("event_normalizer", "expected_frames"), + [ + (None, ["response.cancel", "conversation.item.delete", "conversation.item.create", "response.create"]), + (XAIRealtimeNormalizer(), ["response.cancel", "conversation.item.create", "response.create"]), + ], + ids=["deletes_the_blocked_item", "xai_keeps_the_item_it_cannot_delete"], +) +async def test_blocked_transcript_is_removed_from_the_conversation_before_the_refusal_turn( + monkeypatch: pytest.MonkeyPatch, + event_normalizer: XAIRealtimeNormalizer | None, + expected_frames: list[str], +): + class _BlockingGuardrail(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + raise ValueError("blocked") + + monkeypatch.setattr( + litellm, + "callbacks", + [ + _BlockingGuardrail( + guardrail_name="blocking", event_hook=GuardrailEventHooks.realtime_input_transcription, default_on=True + ) + ], + ) + client_ws: Final = MagicMock() + client_ws.send_text = AsyncMock() + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + backend_frames: Final = iter((_make_transcript_event("blocked words", item_id="item_blocked"),)) + + async def recv(**_kwargs: object) -> bytes: + frame = next(backend_frames, None) + if frame is None: + raise ConnectionClosed(None, None) + return frame + + backend_ws.recv = recv + logging_obj: Final = _mock_logging() + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + + await RealTimeStreaming( + client_ws, backend_ws, logging_obj, event_normalizer=event_normalizer + ).backend_to_client_send_messages() + + sent: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] + assert [event["type"] for event in sent] == expected_frames + assert all(event["item_id"] == "item_blocked" for event in sent if event["type"] == "conversation.item.delete") diff --git a/tests/unit/llms/gemini/realtime/test_gemini_realtime_transformation.py b/tests/unit/llms/gemini/realtime/test_gemini_realtime_transformation.py index bcd5f3d8d19..017f279d6c4 100644 --- a/tests/unit/llms/gemini/realtime/test_gemini_realtime_transformation.py +++ b/tests/unit/llms/gemini/realtime/test_gemini_realtime_transformation.py @@ -1,6 +1,6 @@ import json from collections.abc import Mapping -from typing import cast +from typing import Final, cast from unittest.mock import MagicMock import pytest @@ -1323,6 +1323,40 @@ def test_gemini_realtime_pipecat_semantic_vad_omits_realtime_input_config(): assert setup["tools"][0]["function_declarations"][0]["name"] == "terminate_call" +@pytest.mark.parametrize( + ("create_response", "expected_detection"), + [ + (False, {"disabled": True}), + (True, {"disabled": False, "prefixPaddingMs": 200, "silenceDurationMs": 700}), + ], +) +def test_gemini_turn_detection_tuning_is_dropped_when_auto_response_is_off( + create_response: bool, expected_detection: dict[str, object] +): + session_update: Final = { + "type": "session.update", + "session": { + "audio": { + "input": { + "turn_detection": { + "type": "server_vad", + "create_response": create_response, + "prefix_padding_ms": 200, + "silence_duration_ms": 700, + } + }, + }, + }, + } + messages: Final = GeminiRealtimeConfig().transform_realtime_request( + json.dumps(session_update), + "gemini-live-2.5-flash-native-audio", + session_configuration_request=None, + ) + setup: Final = json.loads(messages[0])["setup"] + assert setup["realtimeInputConfig"]["automaticActivityDetection"] == expected_detection + + def test_gemini_input_audio_buffer_commit_maps_to_audio_stream_end(): config = GeminiRealtimeConfig() setup = { diff --git a/tests/unit/llms/meta/realtime/test_meta_realtime_transformation.py b/tests/unit/llms/meta/realtime/test_meta_realtime_transformation.py index a5d7e47fb65..8b46ea53713 100644 --- a/tests/unit/llms/meta/realtime/test_meta_realtime_transformation.py +++ b/tests/unit/llms/meta/realtime/test_meta_realtime_transformation.py @@ -20,6 +20,8 @@ from litellm.llms.meta.realtime.transformation import ( session_created_event, ) from litellm.types.llms.meta import MuseMode +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager from litellm.types.realtime import RealtimeResponseTransformInput EMPTY_TRANSFORM_INPUT: Final[RealtimeResponseTransformInput] = { @@ -164,6 +166,22 @@ def test_session_rejects_unsupported_audio_model_and_hints(session: dict[str, ob parse_session_update(_event("session.update", session={"type": "transcription", **session}), MUSE_MODEL) +def _handshake(turn_detection: str, drop_params: bool) -> tuple[str | bytes, ...]: + config = ProviderConfigManager.get_provider_realtime_config(MUSE_MODEL, LlmProviders.META, drop_params=drop_params) + assert config is not None + config.validate_environment({}, MUSE_MODEL, api_key="secret-token") + return tuple(config.transform_realtime_request(_ga_session_update(turn_detection=turn_detection), MUSE_MODEL)) + + +def test_semantic_vad_falls_back_to_server_vad_under_drop_params(): + assert _handshake("semantic_vad", drop_params=True) == _handshake("server_vad", drop_params=False) + + +def test_semantic_vad_without_drop_params_names_the_drop_params_fallback(): + with pytest.raises(MuseProtocolError, match="Set drop_params to use server_vad instead"): + _handshake("semantic_vad", drop_params=False) + + def test_session_created_event_exposes_openai_transcription_shape(): config = parse_session_update( _event( diff --git a/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py b/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py index b50d6e19458..5c7099832ba 100644 --- a/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py +++ b/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py @@ -17,6 +17,7 @@ from litellm.llms.vertex_ai.audio_transcription.realtime_transformation import ( parse_chirp_session_update, ) from litellm.llms.vertex_ai.common_utils import VertexAIError +from litellm.llms.vertex_ai.realtime.transformation import vertex_realtime_config from litellm.types.llms.vertex_ai_speech_to_text import ( VertexSpeechStreamingConfigured, VertexSpeechStreamingResponse, @@ -173,6 +174,26 @@ def test_unsupported_session_settings_are_rejected(payload: str, message: str): parse_chirp_session_update(payload, MODEL) +def _configure_command(turn_detection: str, drop_params: bool) -> list[object]: + config = vertex_realtime_config( + MODEL, access_token="token", resolve_access_token=_token, project="proj-1", location="us", drop_params=drop_params + ) + config.transform_session_created_event(MODEL, "sess_1") + return [ + json.loads(command) if isinstance(command, str) else command + for command in config.transform_realtime_request(_ga_session_update(turn_detection=turn_detection), MODEL) + ] + + +def test_semantic_vad_falls_back_to_server_vad_under_drop_params(): + assert _configure_command("semantic_vad", drop_params=True) == _configure_command("server_vad", drop_params=False) + + +def test_semantic_vad_without_drop_params_names_the_drop_params_fallback(): + with pytest.raises(ChirpProtocolError, match="Set drop_params to use server_vad instead"): + _configure_command("semantic_vad", drop_params=False) + + def test_session_update_configures_once_and_later_updates_are_ignored(): config = _config() config.transform_session_created_event(MODEL, "sess_1") diff --git a/tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py b/tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py index a705c169782..0bfff870d15 100644 --- a/tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py +++ b/tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py @@ -223,6 +223,36 @@ def test_vertex_session_update_normalizes_ga_remapped_fields( ) +def test_vertex_keeps_turn_detection_tuning_when_auto_response_is_off(): + session_update: Final = { + "type": "session.update", + "session": { + "audio": { + "input": { + "turn_detection": { + "type": "server_vad", + "create_response": False, + "prefix_padding_ms": 200, + "silence_duration_ms": 700, + } + }, + }, + }, + } + cfg: Final = VertexAIRealtimeConfig(access_token="tok", project="my-proj", location="us-central1") + messages: Final = cfg.transform_realtime_request( + json.dumps(session_update), + _NATIVE_AUDIO_MODEL, + session_configuration_request=None, + ) + setup: Final = json.loads(messages[0])["setup"] + assert setup["realtimeInputConfig"]["automaticActivityDetection"] == { + "disabled": True, + "prefixPaddingMs": 200, + "silenceDurationMs": 700, + } + + # --------------------------------------------------------------------------- # Round-trip test: text-in / text-out via RealTimeStreaming # --------------------------------------------------------------------------- From 977a58121372025f63e74657a589e18e066f59a0 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 16 Oct 2026 08:52:42 -0700 Subject: [PATCH 02/58] fix(router): stop queueing budget increments for Redis when no Redis is configured (#46831) * fix(router): stop queueing budget increments for Redis when no Redis is configured Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): assert Redis-less model budget limiter tracks spend locally without queueing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router_strategy/budget_limiter.py | 5 +- .../sdk/test_router_budget_limiter_wire.py | 94 +++++++++++++++++++ ...test_unit_test_max_model_budget_limiter.py | 8 +- .../router_strategy/test_budget_limiter.py | 38 ++++---- 4 files changed, 119 insertions(+), 26 deletions(-) diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index b20f71cd51f..1622a30b95f 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -405,7 +405,8 @@ class RouterBudgetLimiting(CustomLogger): Runs once the budget start time exists in Redis Cache (on the 2nd and subsequent requests to the same provider) - Increments the spend in memory cache (so spend instantly updated in memory) - - Queues the increment operation to Redis Pipeline (using batched pipeline to optimize performance. Using Redis for multi instance environment of LiteLLM) + - Queues the increment for the batched Redis pipeline when Redis is configured, since only a multi instance + deployment has anything to sync """ increment_op: Final = RedisPipelineIncrementOperation( key=spend_key, @@ -418,6 +419,8 @@ class RouterBudgetLimiting(CustomLogger): value=response_cost, ttl=ttl, ) + if self.dual_cache.redis_cache is None: + return self.redis_increment_operation_queue.append(increment_op) def _get_redis_increment_queue_lock(self) -> asyncio.Lock: diff --git a/tests/integration/sdk/test_router_budget_limiter_wire.py b/tests/integration/sdk/test_router_budget_limiter_wire.py index cf4fda41891..0863ee6d6cc 100644 --- a/tests/integration/sdk/test_router_budget_limiter_wire.py +++ b/tests/integration/sdk/test_router_budget_limiter_wire.py @@ -1,7 +1,9 @@ from __future__ import annotations +import asyncio import json import os +import time import uuid from collections.abc import Iterator from typing import Final @@ -222,3 +224,95 @@ async def test_tag_budget_blocks_exhausted_tag_but_allows_another( assert response.choices[0].message.content == "openai" assert len(requests) == 1 assert json.loads(requests[0].body)["model"] == _MODEL + + + +_ISOLATED_SPEND_KEY: Final = "provider_spend:hosted_vllm:1d" +_INPUT_COST_PER_TOKEN: Final = 0.001 +_OUTPUT_COST_PER_TOKEN: Final = 0.002 + + +def _isolated_provider_deployment(api_base: str) -> dict[str, JsonValue]: + """Only these two tests route to hosted_vllm, so success callbacks still draining from earlier tests cannot land here.""" + return { + "model_name": "isolated-budget", + "litellm_params": { + "model": "hosted_vllm/llama-3", + "api_base": api_base, + "api_key": _API_KEY, + "input_cost_per_token": _INPUT_COST_PER_TOKEN, + "output_cost_per_token": _OUTPUT_COST_PER_TOKEN, + }, + "model_info": {"id": "vllm-deployment"}, + } + + +def _isolated_budget_router(api_base: str, redis_client: Redis | None) -> Router: + redis_params: Final = ( + {} + if redis_client is None + else {"redis_host": os.environ["REDIS_HOST"], "redis_port": int(os.environ["REDIS_PORT"])} + ) + return Router( + model_list=[_isolated_provider_deployment(api_base)], + provider_budget_config={"hosted_vllm": BudgetConfig(budget_duration="1d", max_budget=10.0)}, + num_retries=0, + disable_cooldowns=True, + **redis_params, + ) + + +async def _complete_and_wait_for_spend(router: Router, previous: float) -> float: + response: Final = await router.acompletion( + model="isolated-budget", + messages=[{"role": "user", "content": "budget tracking"}], + ) + cost: Final = ( + response.usage.prompt_tokens * _INPUT_COST_PER_TOKEN + response.usage.completion_tokens * _OUTPUT_COST_PER_TOKEN + ) + deadline: Final = time.monotonic() + 5.0 + while time.monotonic() < deadline: + spend: float | None = router.cache.in_memory_cache.get_cache(_ISOLATED_SPEND_KEY) + if spend is not None and spend > previous: + assert spend == pytest.approx(previous + cost) + return spend + await asyncio.sleep(0.05) + raise AssertionError(f"{_ISOLATED_SPEND_KEY} never rose above {previous}") + + +@pytest.mark.asyncio +async def test_router_without_redis_tracks_spend_but_queues_no_redis_increments() -> None: + """A Router with no Redis never flushes, so anything it queues for Redis is retained for the life of the process.""" + with wire_server(_peer) as wire: + router: Final = _isolated_budget_router(f"{wire.url}/openai/v1", redis_client=None) + first: Final = await _complete_and_wait_for_spend(router, previous=0.0) + second: Final = await _complete_and_wait_for_spend(router, previous=first) + third: Final = await _complete_and_wait_for_spend(router, previous=second) + wire.drain() + + assert third > second > first > 0.0 + assert router.router_budget_logger is not None + assert router.router_budget_logger.redis_increment_operation_queue == [] + + +async def _wait_for_redis_spend(redis_client: Redis, at_least: float) -> float: + deadline: Final = time.monotonic() + 15.0 + while time.monotonic() < deadline: + raw: str | None = redis_client.get(_ISOLATED_SPEND_KEY) + if raw is not None and float(raw) >= at_least: + return float(raw) + await asyncio.sleep(0.1) + raise AssertionError(f"{_ISOLATED_SPEND_KEY} in Redis never reached {at_least}") + + +@pytest.mark.asyncio +async def test_router_with_redis_flushes_queued_spend_increments(redis_client: Redis) -> None: + redis_client.delete(_ISOLATED_SPEND_KEY) + with wire_server(_peer) as wire: + router: Final = _isolated_budget_router(f"{wire.url}/openai/v1", redis_client=redis_client) + first: Final = await _complete_and_wait_for_spend(router, previous=0.0) + second: Final = await _complete_and_wait_for_spend(router, previous=first) + redis_spend: Final = await _wait_for_redis_spend(redis_client, at_least=second) + wire.drain() + + assert redis_spend == pytest.approx(second) diff --git a/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py b/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py index 33904637366..c6369a1f3d7 100644 --- a/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py +++ b/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py @@ -20,7 +20,6 @@ from litellm.proxy.hooks.model_max_budget_limiter import ( model_budget_start_time_cache_key, resolve_model_budget, ) -from litellm.types.caching import RedisPipelineIncrementOperation async def _reset_at_of_open_window( @@ -500,16 +499,15 @@ async def test_async_log_success_event_pushes_redis_increments_when_redis_config @pytest.mark.asyncio -async def test_model_budget_limiter_initializes_redis_increment_queue_lock(): +async def test_model_budget_limiter_without_redis_tracks_spend_locally_and_queues_nothing(): dual_cache = DualCache() limiter = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) spend_key = "virtual_key_spend:test-key:gpt-4:1d" await limiter._increment_spend_in_current_window(spend_key=spend_key, response_cost=0.01, ttl=86400) - assert limiter.redis_increment_operation_queue == [ - RedisPipelineIncrementOperation(key=spend_key, increment_value=0.01, ttl=86400) - ] + assert await dual_cache.async_get_cache(spend_key) == 0.01 + assert limiter.redis_increment_operation_queue == [] @pytest.mark.asyncio diff --git a/tests/unit/router_strategy/test_budget_limiter.py b/tests/unit/router_strategy/test_budget_limiter.py index 685ea46547e..380aa573b3c 100644 --- a/tests/unit/router_strategy/test_budget_limiter.py +++ b/tests/unit/router_strategy/test_budget_limiter.py @@ -14,6 +14,7 @@ import pytest from litellm import Router from litellm.caching.caching import DualCache +from litellm.caching.in_memory_cache import InMemoryCache from litellm.router_strategy.budget_limiter import RouterBudgetLimiting from litellm.types.utils import BudgetConfig import sys, os, asyncio, time, random @@ -232,7 +233,11 @@ async def test_a_monthly_router_budget_ends_a_calendar_month_after_it_opens( """ monkeypatch.setattr(litellm, "tag_budget_config", tag_budget_config) monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) - limiter = RouterBudgetLimiting(dual_cache=DualCache(), clock=charged_at.timestamp, **limiter_config) + limiter = RouterBudgetLimiting( + dual_cache=DualCache(in_memory_cache=InMemoryCache(clock=charged_at.timestamp)), + clock=charged_at.timestamp, + **limiter_config, + ) await limiter.dual_cache.async_set_cache(key=start_time_key, value=opened_at.timestamp()) await _log_success( @@ -243,9 +248,8 @@ async def test_a_monthly_router_budget_ends_a_calendar_month_after_it_opens( }, ) - assert limiter.redis_increment_operation_queue == [ - {"key": spend_key, "increment_value": 0.25, "ttl": int((closes_at - charged_at).total_seconds())} - ] + assert await limiter.dual_cache.async_get_cache(spend_key) == 0.25 + assert limiter.dual_cache.in_memory_cache.ttl_dict[spend_key] == closes_at.timestamp() @pytest.mark.asyncio @@ -362,24 +366,18 @@ async def test_new_budget_window_sets_spend_and_start_time() -> None: @pytest.mark.asyncio -async def test_spend_increment_updates_memory_and_queues_redis_operation() -> None: - provider_budget: Final = RouterBudgetLimiting(dual_cache=DualCache(), provider_budget_config={}) - await provider_budget.dual_cache.async_set_cache(key="provider_spend:openai:1d", value=1.0, ttl=86400) - - await provider_budget._increment_spend_in_current_window( - spend_key="provider_spend:openai:1d", - response_cost=0.5, - ttl=86400, +async def test_redis_less_limiter_tracks_spend_without_queueing_redis_increments(disable_budget_sync) -> None: + """With no Redis there is nothing to flush to, so queued increments would be retained for the life of the process.""" + limiter: Final = RouterBudgetLimiting( + dual_cache=DualCache(), + provider_budget_config={"openai": {"budget_limit": 10.0, "time_period": "1d"}}, ) - assert await provider_budget.dual_cache.async_get_cache("provider_spend:openai:1d") == 1.5 - assert provider_budget.redis_increment_operation_queue == [ - { - "key": "provider_spend:openai:1d", - "increment_value": 0.5, - "ttl": 86400, - } - ] + for _ in range(3): + await _log_success(limiter, _success_kwargs(provider_in_litellm_params="openai", provider_in_payload="openai")) + + assert await limiter.dual_cache.async_get_cache("provider_spend:openai:1d") == pytest.approx(0.75) + assert limiter.redis_increment_operation_queue == [] @pytest.mark.asyncio From e45f2737d33772efd72df232cdc7a6e56154dfd7 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 16 Oct 2026 09:22:22 -0700 Subject: [PATCH 03/58] fix(proxy)!: delete callbacks configured under callbacks or failure_callback (#42740) * fix(proxy): delete callbacks configured under callbacks or failure_callback Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(e2e): assert on configured callback rows and use a callback absent from CI configs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(e2e): follow logging suite conventions for callback delete tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): flatten remaining callbacks with a helper generator Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): type callback delete helpers to satisfy basedpyright budget Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): save callback config without mutating the loaded dict Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): delete callbacks from every litellm_settings list GET shows POST /config/callback/delete only searched success_callback while GET /get/config/callbacks also lists failure_callback and callbacks, so the Admin UI Delete returned 404 for those rows Remove the name from every list that holds it, matching entries case-insensitively, and skip lists the config file owns when the database owns another copy so the save is not refused Replace the e2e test with an integration test: the shared e2e stack config owns litellm_settings.callbacks, so it cannot exercise that list * test(proxy): cover config-owned success_callback on callback delete and wait on runtime state Add the mirror mixed-ownership case where the config file sets success_callback and the database owns failure_callback, which main refused with a 400, plus a string-valued list case Poll GET listings and /active/callbacks instead of reading them once, so a request served by a worker that has not synced yet cannot fail the test, and reuse the existing callback list key map in the route * test(proxy): annotate callback delete test helper locals with Final * fix(proxy): save callback delete through the config it read so concurrent settings edits are kept --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 62 +++- .../configuration/test_callback_delete.py | 351 ++++++++++++++++++ .../proxy/proxy_server/test_routes_config.py | 227 +++++++++++ 3 files changed, 628 insertions(+), 12 deletions(-) create mode 100644 tests/integration/configuration/test_callback_delete.py diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9eb2bc93391..b56f0e68c4f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -19698,6 +19698,26 @@ async def delete_config_general_settings( return response +def _configured_callback_names(value: object) -> tuple[JsonValue, ...]: + if isinstance(value, str): + return (value,) + if isinstance(value, list): + return tuple(cast(Iterable[JsonValue], value)) # cast-ok: config.yaml values are untyped + return () + + +def _is_callback_name(entry: object, callback_name: str) -> bool: + return isinstance(entry, str) and entry.lower() == callback_name + + +def _callback_lists_to_rewrite( + matching_keys: tuple[str, ...], owned_by_config: Callable[[str], bool] +) -> tuple[str, ...]: + """The matching lists the database owns, or all of them when the config file owns every one.""" + database_owned: Final = tuple(key for key in matching_keys if not owned_by_config(key)) + return database_owned or matching_keys + + @router.post( "/config/callback/delete", tags=["config.yaml"], @@ -19733,24 +19753,42 @@ async def delete_callback( try: # Get current configuration - config: Final = await proxy_config.get_config() + config: Final = cast( # cast-ok: get_config returns an untyped config dict + dict[str, object], + await proxy_config.get_config(), + ) callback_name: Final = data.callback_name.lower() # Check if callback exists in current configuration - litellm_settings: Final = config.get("litellm_settings", {}) - success_callbacks: Final = litellm_settings.get("success_callback", []) + litellm_settings: Final[Mapping[str, object]] = cast( # cast-ok: the config dict's values are untyped + Mapping[str, object], + config.get("litellm_settings", {}), + ) + configured_lists: Final = { + key: _configured_callback_names(litellm_settings.get(key)) for key in _DB_CONFIG_CALLBACK_EVENT_TYPES + } + matching_keys: Final = tuple( + key + for key, names in configured_lists.items() + if any(_is_callback_name(entry, callback_name) for entry in names) + ) - if callback_name not in success_callbacks: + if not matching_keys: raise HTTPException( status_code=404, detail={"error": f"Callback '{callback_name}' not found in active configuration"}, ) - before_success_callbacks: Final = list(success_callbacks) - - # Remove callback from success_callback list - success_callbacks.remove(callback_name) - config.setdefault("litellm_settings", {})["success_callback"] = success_callbacks + rewritten_keys: Final = _callback_lists_to_rewrite(matching_keys, proxy_config.litellm_settings.owned_by_config) + remaining_lists: Final = MappingProxyType( + { + key: tuple(entry for entry in configured_lists[key] if not _is_callback_name(entry, callback_name)) + for key in rewritten_keys + } + ) + before_callbacks: Final[JsonValue] = {key: list(configured_lists[key]) for key in rewritten_keys} + after_callbacks: Final[JsonValue] = {key: list(names) for key, names in remaining_lists.items()} + config["litellm_settings"] = {**litellm_settings, **after_callbacks} # Save the updated configuration await proxy_config.save_config(new_config=config) @@ -19759,8 +19797,8 @@ async def delete_callback( create_config_audit_log( "litellm_settings", "deleted", - {"success_callback": before_success_callbacks}, - {"success_callback": success_callbacks}, + before_callbacks, + after_callbacks, user_api_key_dict, ) ) @@ -19771,7 +19809,7 @@ async def delete_callback( return { "message": f"Successfully deleted callback: {callback_name}", "removed_callback": callback_name, - "remaining_callbacks": success_callbacks, + "remaining_callbacks": tuple(chain.from_iterable(remaining_lists.values())), "deleted_at": datetime.now().isoformat(), } diff --git a/tests/integration/configuration/test_callback_delete.py b/tests/integration/configuration/test_callback_delete.py new file mode 100644 index 00000000000..d58dd252593 --- /dev/null +++ b/tests/integration/configuration/test_callback_delete.py @@ -0,0 +1,351 @@ +import json +import os +import uuid +from collections.abc import Generator, Iterator, Mapping +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import pytest +import yaml +from integration._support.client import JSON_OBJECT, Gateway, JsonValue, eventually, list_value, object_value +from integration._support.database import read_rows, scratch_database +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, Wire, wire_server + +CALLBACK: Final = "generic_api" +SIBLING: Final = "langsmith" +LIST_TYPES: Final = MappingProxyType( + { + "success_callback": "success", + "failure_callback": "failure", + "callbacks": "success_and_failure", + } +) +SINK_ENVIRONMENT: Final = MappingProxyType( + { + "GENERIC_LOGGER_ENDPOINT": "http://127.0.0.1:9/unused", + "LANGSMITH_API_KEY": "synthetic-langsmith-key", + "LANGSMITH_BASE_URL": "http://127.0.0.1:9", + } +) +NOT_FOUND: Final = f"Callback '{CALLBACK}' not found in active configuration" +WORKER_SYNC_SECONDS: Final = 40 + + +@dataclass(frozen=True, slots=True) +class Rig: + proxy: Gateway + database_url: str + + +@contextmanager +def _rig(gateway: Gateway, tmp_path: Path, file_settings: Mapping[str, JsonValue]) -> Generator[Rig]: + config: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + path: Final = tmp_path / f"callback_delete_{uuid.uuid4().hex}.yaml" + path.write_text( + yaml.safe_dump({**config, "litellm_settings": {**object_value(config["litellm_settings"]), **file_settings}}) + ) + with ( + scratch_database() as database_url, + owned_proxy( + gateway, + tmp_path, + {"DATABASE_URL": database_url, **SINK_ENVIRONMENT}, + config=path, + remove_environment=("DATABASE_URL_READ_REPLICA",), + workers=2, + ) as proxy, + ): + yield Rig(proxy, database_url) + + +@pytest.fixture +def rig(gateway: Gateway, tmp_path: Path) -> Iterator[Rig]: + with _rig(gateway, tmp_path, {}) as created: + yield created + + +def _configured(proxy: Gateway) -> tuple[tuple[str, str], ...]: + rows: Final = proxy.get("/get/config/callbacks")["callbacks"] + assert isinstance(rows, list) + return tuple( + sorted( + (str(row["name"]), str(row["type"])) + for row in (object_value(entry) for entry in rows) + if not row.get("read_only") + ) + ) + + +def _listed(proxy: Gateway, expected: tuple[tuple[str, str], ...]) -> None: + eventually(lambda: _configured(proxy), lambda rows: rows == expected, seconds=WORKER_SYNC_SECONDS) + + +def _stored_settings(database_url: str) -> Mapping[str, JsonValue]: + rows: Final = read_rows( + 'SELECT param_value FROM "LiteLLM_Config" WHERE param_name = %s', + ("litellm_settings",), + database_url=database_url, + ) + assert len(rows) == 1, rows + value: Final = rows[0]["param_value"] + return object_value(json.loads(value) if isinstance(value, str) else value) + + +def _configure( + proxy: Gateway, settings: Mapping[str, JsonValue], environment: Mapping[str, JsonValue] | None = None +) -> None: + body: Final[Mapping[str, JsonValue]] = { + "litellm_settings": dict(settings), + "environment_variables": None if environment is None else dict(environment), + } + assert proxy.post("/config/update", body) == {"message": "Config updated successfully"} + + +def _delete(proxy: Gateway, name: str) -> tuple[int, Mapping[str, JsonValue]]: + response: Final = proxy.request("POST", "/config/callback/delete", {"callback_name": name}) + return response.status_code, JSON_OBJECT.validate_json(response.content) + + +@pytest.mark.parametrize("list_key", tuple(LIST_TYPES)) +def test_callback_listed_under_any_settings_list_is_deleted_from_the_listing_and_the_store( + rig: Rig, list_key: str +) -> None: + _configure(rig.proxy, {list_key: [CALLBACK]}) + _listed(rig.proxy, ((CALLBACK, LIST_TYPES[list_key]),)) + + status, body = _delete(rig.proxy, CALLBACK) + + assert status == 200, body + assert (body["removed_callback"], body["remaining_callbacks"]) == (CALLBACK, []) + assert _stored_settings(rig.database_url)[list_key] == [] + _listed(rig.proxy, ()) + + +def test_callback_in_both_success_and_failure_lists_is_removed_from_both_in_one_delete(rig: Rig) -> None: + _configure(rig.proxy, {"success_callback": [CALLBACK], "failure_callback": [CALLBACK]}) + _listed(rig.proxy, ((CALLBACK, "failure"), (CALLBACK, "success"))) + + status, body = _delete(rig.proxy, CALLBACK) + + assert status == 200, body + stored: Final = _stored_settings(rig.database_url) + assert (stored["success_callback"], stored["failure_callback"]) == ([], []) + _listed(rig.proxy, ()) + + +def test_delete_keeps_other_callbacks_and_reports_what_is_left_in_the_list_it_edited(rig: Rig) -> None: + _configure(rig.proxy, {"callbacks": [SIBLING, CALLBACK], "failure_callback": [SIBLING]}) + _listed(rig.proxy, ((CALLBACK, "success_and_failure"), (SIBLING, "failure"), (SIBLING, "success_and_failure"))) + + status, body = _delete(rig.proxy, CALLBACK) + + assert status == 200, body + assert body["remaining_callbacks"] == [SIBLING] + stored: Final = _stored_settings(rig.database_url) + assert (stored["callbacks"], stored["failure_callback"]) == ([SIBLING], [SIBLING]) + _listed(rig.proxy, ((SIBLING, "failure"), (SIBLING, "success_and_failure"))) + + +def test_delete_matches_the_callback_name_case_insensitively(rig: Rig) -> None: + _configure(rig.proxy, {"failure_callback": [CALLBACK.upper()], "callbacks": [CALLBACK]}) + _listed(rig.proxy, ((CALLBACK.upper(), "failure"), (CALLBACK, "success_and_failure"))) + + status, body = _delete(rig.proxy, CALLBACK.upper()) + + assert status == 200, body + assert body["removed_callback"] == CALLBACK + _listed(rig.proxy, ()) + + +def test_delete_of_an_unconfigured_callback_is_a_404_and_changes_nothing(rig: Rig) -> None: + _configure(rig.proxy, {"callbacks": [SIBLING]}) + _listed(rig.proxy, ((SIBLING, "success_and_failure"),)) + before: Final = _stored_settings(rig.database_url) + + status, body = _delete(rig.proxy, CALLBACK) + + assert (status, body) == (404, {"detail": {"error": NOT_FOUND}}) + assert _stored_settings(rig.database_url) == before + _listed(rig.proxy, ((SIBLING, "success_and_failure"),)) + + +@pytest.mark.parametrize( + ("file_key", "database_key"), (("failure_callback", "success_callback"), ("success_callback", "failure_callback")) +) +def test_delete_removes_the_database_copy_and_keeps_the_copy_the_config_file_sets( + gateway: Gateway, tmp_path: Path, file_key: str, database_key: str +) -> None: + with _rig(gateway, tmp_path, {file_key: [CALLBACK]}) as file_owned: + _configure(file_owned.proxy, {database_key: [CALLBACK]}) + _listed(file_owned.proxy, ((CALLBACK, "failure"), (CALLBACK, "success"))) + + status, body = _delete(file_owned.proxy, CALLBACK) + + assert status == 200, body + assert body["remaining_callbacks"] == [] + assert _stored_settings(file_owned.database_url)[database_key] == [] + _listed(file_owned.proxy, ((CALLBACK, LIST_TYPES[file_key]),)) + + +def test_delete_of_a_callback_only_the_config_file_lists_is_refused_and_names_the_file( + gateway: Gateway, tmp_path: Path +) -> None: + with _rig(gateway, tmp_path, {"callbacks": [CALLBACK]}) as file_owned: + _listed(file_owned.proxy, ((CALLBACK, "success_and_failure"),)) + + status, body = _delete(file_owned.proxy, CALLBACK) + + assert status == 400, body + detail: Final = object_value(body["detail"]) + assert (detail["keys"], detail["section"]) == (["callbacks"], "litellm_settings") + assert "set in the config file" in str(detail["error"]) + _listed(file_owned.proxy, ((CALLBACK, "success_and_failure"),)) + + +def _completion(request: Request) -> Reply: + message: Final = object_value(list_value(JSON_OBJECT.validate_json(request.body)["messages"])[0]) + if str(message["content"]).endswith("-failure"): + return Reply( + status=400, + body=json.dumps( + {"error": {"type": "invalid_request_error", "code": "synthetic", "message": "synthetic failure"}} + ).encode(), + ) + return Reply( + body=json.dumps( + { + "id": "chatcmpl-" + uuid.uuid4().hex, + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}, + } + ).encode() + ) + + +@dataclass(slots=True) +class Sink: + wire: Wire + bodies: list[bytes] # mutable-ok: drain() consumes the queue, later reads keep earlier batches + + def received(self, markers: tuple[str, ...]) -> bool: + self.bodies.extend(request.body for request in self.wire.drain()) + return all(any(marker.encode() in body for body in self.bodies) for marker in markers) + + +@dataclass(frozen=True, slots=True) +class Pair: + first: Gateway + peer: Gateway + database_url: str + provider: Wire + generic: Sink + langsmith: Sink + + def running(self, name: str) -> tuple[bool, bool]: + return (_runs(self.first, name), _runs(self.peer, name)) + + +def _runs(proxy: Gateway, name: str) -> bool: + active: Final = list_value(proxy.get("/active/callbacks")["all_litellm_callbacks"]) + return any(name in str(entry).lower() for entry in active) + + +@contextmanager +def _pair(gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Generator[Pair]: + with ( + scratch_database() as database_url, + wire_server(_completion) as provider, + wire_server(lambda _: Reply()) as generic, + wire_server(lambda _: Reply()) as langsmith, + ): + monkeypatch.setenv("DATABASE_URL", database_url) + environment: Final = { + "DATABASE_URL": database_url, + "DEFAULT_FLUSH_INTERVAL_SECONDS": "1", + "LITELLM_STORE_AUDIT_LOGS": "True", + } + removed: Final = ( + "DATABASE_URL_READ_REPLICA", + "GENERIC_LOGGER_ENDPOINT", + "GENERIC_LOGGER_HEADERS", + "LANGSMITH_API_KEY", + "LANGSMITH_BASE_URL", + ) + with ( + owned_proxy(gateway, tmp_path, environment, remove_environment=removed) as first, + owned_proxy(gateway, tmp_path, environment, remove_environment=removed) as peer, + ): + yield Pair(first, peer, database_url, provider, Sink(generic, []), Sink(langsmith, [])) + + +def _send_on_both(pair: Pair, model: str, key: str, marker: str, outcome: str) -> tuple[str, str]: + markers: Final = (f"{marker}-first-{outcome}", f"{marker}-peer-{outcome}") + for proxy, text in zip((pair.first, pair.peer), markers, strict=True): + response = proxy.request( + "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": text}]}, key=key + ) + assert response.status_code == (400 if outcome == "failure" else 200), response.text + return markers + + +@pytest.mark.parametrize(("list_key", "outcome"), (("failure_callback", "failure"), ("callbacks", "success"))) +def test_deleted_callback_stops_receiving_events_on_both_processes( + gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, list_key: str, outcome: str +) -> None: + assert os.environ.get("LITELLM_LICENSE"), "LITELLM_LICENSE must be set: audit logs are an enterprise feature" + marker: Final = "callback-delete-" + uuid.uuid4().hex + with _pair(gateway, tmp_path, monkeypatch) as pair, pair.first.scenario() as scenario: + _configure( + pair.first, + {list_key: [SIBLING, CALLBACK]}, + { + "GENERIC_LOGGER_ENDPOINT": pair.generic.wire.url, + "LANGSMITH_API_KEY": "langsmith-key-" + marker, + "LANGSMITH_BASE_URL": pair.langsmith.wire.url, + }, + ) + eventually(lambda: pair.running(SIBLING), lambda running: running == (True, True), seconds=WORKER_SYNC_SECONDS) + model: Final = scenario.model(api_base=pair.provider.url + "/v1") + key: Final = scenario.key(models=[model]) + before: Final = _send_on_both(pair, model, key, marker + "-before", outcome) + eventually(lambda: pair.generic.received(before), bool, seconds=25) + eventually(lambda: pair.langsmith.received(before), bool, seconds=25) + + status, body = _delete(pair.first, SIBLING) + + assert status == 200, body + assert (body["removed_callback"], body["remaining_callbacks"]) == (SIBLING, [CALLBACK]) + assert _stored_settings(pair.database_url)[list_key] == [CALLBACK] + audit: Final = eventually( + lambda: read_rows( + 'SELECT before_value, updated_values FROM "LiteLLM_AuditLog" WHERE object_id = %s AND action = %s', + ("litellm_settings", "deleted"), + database_url=pair.database_url, + ), + bool, + seconds=10, + ) + assert len(audit) == 1, audit + before_value: Final = list_value(object_value(audit[0]["before_value"])[list_key]) + assert sorted(str(name) for name in before_value) == sorted([SIBLING, CALLBACK]), audit + assert audit[0]["updated_values"] == {list_key: [CALLBACK]}, audit + eventually( + lambda: pair.running(SIBLING), lambda running: running == (False, False), seconds=WORKER_SYNC_SECONDS + ) + + after: Final = _send_on_both(pair, model, key, marker + "-after", outcome) + eventually(lambda: pair.generic.received(after), bool, seconds=25) + _configure(pair.first, {list_key: [CALLBACK, SIBLING]}) + eventually(lambda: pair.running(SIBLING), lambda running: running == (True, True), seconds=WORKER_SYNC_SECONDS) + barrier: Final = _send_on_both(pair, model, key, marker + "-barrier", outcome) + eventually(lambda: pair.langsmith.received(barrier), bool, seconds=25) + + leaked: Final = tuple(text for text in after if any(text.encode() in body for body in pair.langsmith.bodies)) + assert leaked == (), leaked diff --git a/tests/unit/proxy/proxy_server/test_routes_config.py b/tests/unit/proxy/proxy_server/test_routes_config.py index 4c8265b82c3..2d86ff6a7c2 100644 --- a/tests/unit/proxy/proxy_server/test_routes_config.py +++ b/tests/unit/proxy/proxy_server/test_routes_config.py @@ -14,10 +14,15 @@ Routes covered: from __future__ import annotations import asyncio +import contextlib import json +from collections.abc import Callable +from typing import Final from unittest.mock import AsyncMock, MagicMock +import httpx import pytest +from fastapi.testclient import TestClient from .conftest import VOLATILE_KEYS, normalize @@ -944,6 +949,228 @@ def test_config_callback_delete_not_found(client, auth_as, mock_prisma, monkeypa assert "langfuse" in str(response.json()).lower() or "not found" in str(response.json()).lower() +def _delete_callback_roundtrip( + client: TestClient, + auth_as: Callable[..., contextlib.AbstractContextManager[None]], + mock_prisma: MagicMock, + monkeypatch: pytest.MonkeyPatch, + litellm_settings: dict[str, object], + callback_name: str, + owned_by_config: frozenset[str] = frozenset(), +) -> tuple[httpx.Response, dict[str, object]]: + """POST /config/callback/delete as admin against a stubbed config, and + return (response, saved litellm_settings) so each case pins the full body + plus the exact config persisted.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "store_model_in_db", True) + + fake_proxy_config: Final = MagicMock() + fake_proxy_config.get_config = AsyncMock(return_value={"litellm_settings": litellm_settings}) + fake_proxy_config.save_config = AsyncMock() + fake_proxy_config.add_deployment = AsyncMock() + fake_proxy_config.litellm_settings.owned_by_config = owned_by_config.__contains__ + monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response: Final = client.post("/config/callback/delete", json={"callback_name": callback_name}) + saved: Final = fake_proxy_config.save_config.await_args.kwargs["new_config"]["litellm_settings"] + return response, saved + + +def test_config_callback_delete_from_failure_callback( + client: TestClient, + auth_as: Callable[..., contextlib.AbstractContextManager[None]], + mock_prisma: MagicMock, + monkeypatch: pytest.MonkeyPatch, +): + """A callback configured only under litellm_settings.failure_callback is + listed by GET /get/config/callbacks, so the delete route must remove it.""" + response, saved = _delete_callback_roundtrip( + client, + auth_as, + mock_prisma, + monkeypatch, + {"success_callback": ["slack"], "failure_callback": ["datadog", "sentry"]}, + "datadog", + ) + assert response.status_code == 200 + assert response.json()["removed_callback"] == "datadog" + assert saved == {"success_callback": ["slack"], "failure_callback": ["sentry"]} + assert response.json()["remaining_callbacks"] == ["sentry"] + + +def test_config_callback_delete_from_callbacks_list( + client: TestClient, + auth_as: Callable[..., contextlib.AbstractContextManager[None]], + mock_prisma: MagicMock, + monkeypatch: pytest.MonkeyPatch, +): + """A callback configured only under litellm_settings.callbacks is deleted + while unrelated lists pass through untouched.""" + response, saved = _delete_callback_roundtrip( + client, + auth_as, + mock_prisma, + monkeypatch, + {"success_callback": ["langfuse"], "callbacks": ["datadog"], "drop_params": True}, + "datadog", + ) + assert response.status_code == 200 + assert response.json()["removed_callback"] == "datadog" + assert saved == {"success_callback": ["langfuse"], "callbacks": [], "drop_params": True} + assert response.json()["remaining_callbacks"] == [] + + +def test_config_callback_delete_from_multiple_lists( + client: TestClient, + auth_as: Callable[..., contextlib.AbstractContextManager[None]], + mock_prisma: MagicMock, + monkeypatch: pytest.MonkeyPatch, +): + """A name present in both success_callback and callbacks is removed from + both, and remaining_callbacks reports the leftovers of the lists it left.""" + response, saved = _delete_callback_roundtrip( + client, + auth_as, + mock_prisma, + monkeypatch, + { + "success_callback": ["datadog", "slack"], + "failure_callback": ["sentry"], + "callbacks": ["langfuse", "datadog"], + }, + "datadog", + ) + assert response.status_code == 200 + assert saved == { + "success_callback": ["slack"], + "failure_callback": ["sentry"], + "callbacks": ["langfuse"], + } + assert response.json()["remaining_callbacks"] == ["slack", "langfuse"] + + +def test_config_callback_delete_skips_lists_the_config_file_owns( + client: TestClient, + auth_as: Callable[..., contextlib.AbstractContextManager[None]], + mock_prisma: MagicMock, + monkeypatch: pytest.MonkeyPatch, +): + """A name in a database-owned list and a config-file-owned list is removed + from the database-owned one only, so the save is not refused.""" + response, saved = _delete_callback_roundtrip( + client, + auth_as, + mock_prisma, + monkeypatch, + {"success_callback": ["datadog", "slack"], "failure_callback": ["datadog"]}, + "datadog", + owned_by_config=frozenset({"failure_callback"}), + ) + assert response.status_code == 200 + assert saved == {"success_callback": ["slack"], "failure_callback": ["datadog"]} + assert response.json()["remaining_callbacks"] == ["slack"] + + +def test_config_callback_delete_removes_database_copy_when_config_file_owns_success_callback( + client: TestClient, + auth_as: Callable[..., contextlib.AbstractContextManager[None]], + mock_prisma: MagicMock, + monkeypatch: pytest.MonkeyPatch, +): + """A name the config file sets in success_callback is left there while its + database-owned failure_callback copy is removed.""" + response, saved = _delete_callback_roundtrip( + client, + auth_as, + mock_prisma, + monkeypatch, + {"success_callback": ["datadog"], "failure_callback": ["datadog", "sentry"]}, + "datadog", + owned_by_config=frozenset({"success_callback"}), + ) + assert response.status_code == 200 + assert saved == {"success_callback": ["datadog"], "failure_callback": ["sentry"]} + assert response.json()["remaining_callbacks"] == ["sentry"] + + +def test_config_callback_delete_from_string_valued_list( + client: TestClient, + auth_as: Callable[..., contextlib.AbstractContextManager[None]], + mock_prisma: MagicMock, + monkeypatch: pytest.MonkeyPatch, +): + """A list written as a single string is matched and saved back as an empty list.""" + response, saved = _delete_callback_roundtrip( + client, + auth_as, + mock_prisma, + monkeypatch, + {"failure_callback": "datadog"}, + "datadog", + ) + assert response.status_code == 200 + assert saved == {"failure_callback": []} + + +def test_config_callback_delete_case_insensitive( + client: TestClient, + auth_as: Callable[..., contextlib.AbstractContextManager[None]], + mock_prisma: MagicMock, + monkeypatch: pytest.MonkeyPatch, +): + """A mixed-case entry is deleted by the name GET /get/config/callbacks shows for it.""" + response, saved = _delete_callback_roundtrip( + client, + auth_as, + mock_prisma, + monkeypatch, + {"callbacks": ["DataDog", "otel"]}, + "DataDog", + ) + assert response.status_code == 200 + assert response.json()["removed_callback"] == "datadog" + assert saved == {"callbacks": ["otel"]} + + +def test_config_callback_delete_saves_against_the_baseline_it_read( + client: TestClient, + auth_as: Callable[..., contextlib.AbstractContextManager[None]], + mock_prisma: MagicMock, + monkeypatch: pytest.MonkeyPatch, +): + """save_config gets the object get_config returned, so it diffs against the + settings as they were read and writes only the edited callback list.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "store_model_in_db", True) + + read_time: Final = {"litellm_settings": {"failure_callback": ["datadog", "sentry"], "max_budget": 10}} + loaded: Final = ps._ConfigWithBaseline(read_time) + fake_proxy_config: Final = MagicMock() + fake_proxy_config.get_config = AsyncMock(return_value=loaded) + fake_proxy_config.save_config = AsyncMock() + fake_proxy_config.add_deployment = AsyncMock() + fake_proxy_config.litellm_settings.owned_by_config = frozenset[str]().__contains__ + monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response: Final = client.post("/config/callback/delete", json={"callback_name": "datadog"}) + + assert response.status_code == 200 + saved: Final = fake_proxy_config.save_config.await_args.kwargs["new_config"] + assert saved is loaded + assert saved.baseline == read_time + assert saved["litellm_settings"] == {"failure_callback": ["sentry"], "max_budget": 10} + + # --------------------------------------------------------------------------- # GET /get/config/callbacks # --------------------------------------------------------------------------- From e4dc696af85d2034ffd5824460c596f6dd9b2fc1 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 16 Oct 2026 09:33:49 -0700 Subject: [PATCH 04/58] feat(ui): turn the Chat UI on by default (#46918) Co-authored-by: Krrish Dholakia Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy_setting_endpoints.py | 2 +- .../test_proxy_setting_endpoints.py | 18 ++++++++++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index b327129793d..4ece88f757d 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -343,7 +343,7 @@ class UISettings(LiteLLMBaseModel): ) enable_chat_ui: bool = Field( - default=False, + default=True, description="If true, shows the Chat page in the UI sidebar, letting users chat with an LLM and connect their own MCP server credentials via OAuth.", ) diff --git a/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index 6705e5ac8a5..da799a307a6 100644 --- a/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -3771,6 +3771,24 @@ class TestPtuCostAttributionUISetting: assert not mock_prisma.db.litellm_uisettings.upsert.called +class TestEnableChatUISetting: + def test_chat_ui_is_on_when_nothing_is_stored(self, mock_auth, monkeypatch): + TestPtuCostAttributionUISetting._mock_prisma(monkeypatch) + + response = client.get("/get/ui_settings") + + assert response.status_code == 200 + assert response.json()["values"]["enable_chat_ui"] is True + + def test_admin_can_still_turn_chat_ui_off(self, mock_auth, monkeypatch): + TestPtuCostAttributionUISetting._mock_prisma(monkeypatch, stored={"enable_chat_ui": False}) + + response = client.get("/get/ui_settings") + + assert response.status_code == 200 + assert response.json()["values"]["enable_chat_ui"] is False + + class TestApplyUserBudgetToTeamKeysUISetting: """``apply_user_budget_to_team_keys`` mirrors general_settings on every GET. From 7ddc9f1c96c15321a26ef3ae01815d07cad33185 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 17:27:57 +0000 Subject: [PATCH 05/58] fix(bedrock)!: keep reasoning, thinking blocks and the json_mode finish_reason on fake streams (#46863) * fix(bedrock): keep reasoning, thinking blocks and the json_mode finish_reason on fake streams A Bedrock Converse model without native structured output fakes the stream when response_format carries a json_schema. The fake stream transformed the response without json_mode and rebuilt a chunk from text, tool_use and token counts only, so it dropped reasoning_content, the signed thinking_blocks and the reasoning token count, and ended on finish_reason tool_calls. The three fake-stream sites now pass json_mode to _transform_response and stream the resulting ModelResponse through the shared convert_model_response_to_streaming converter, which now numbers each streamed thinking block with its index the way real streams do * fix(bedrock): parse text tool calls on Converse fake streams like the non-stream answer * fix(responses): stream the answer text when one chat chunk carries reasoning and the answer * fix(streaming): keep every field on the OpenAI and openai_like fake streams and split multi-block thinking Azure o-series and openai_like fake streams now use the shared MockResponseIterator, so extra tool calls, reasoning and cached and reasoning tokens survive; the Bedrock-local iterator is deleted. The /v1/responses bridge also splits a chunk whose thinking block is signed or redacted with no reasoning text, and the sync iterator closes the reasoning item before the answer opens, like the async one. The /v1/messages bridge gives each thinking block in a chunk its own content block, starts a new block when the thinking block index changes, and streams redacted thinking as redacted_thinking. Snowflake replays drop the stream index from thinking blocks. * fix(anthropic): stream each thinking block's text once on /v1/messages An unsigned thinking chunk now opens its content block with empty text, since its thinking_delta carries the text. A multi-block chunk keeps each piece's thinking block, so an unsigned block followed by a signed one keeps its own text and index. * docs(anthropic): describe which thinking block start carries its text * fix(streaming): close merged reasoning before the answer when one chunk carries both * test(streaming): cover fake-stream reasoning across the proxy bridges and type the Snowflake replay test --- .../litellm_core_utils/streaming_handler.py | 42 +- .../adapters/streaming_iterator.py | 75 +- .../pass_through/adapters/transformation.py | 8 +- litellm/llms/base_llm/base_model_iterator.py | 26 +- litellm/llms/bedrock/chat/converse_handler.py | 19 +- litellm/llms/bedrock/chat/invoke_handler.py | 125 +-- litellm/llms/openai/openai.py | 3 +- litellm/llms/openai_like/chat/handler.py | 2 +- litellm/llms/snowflake/chat/transformation.py | 8 +- .../streaming_iterator.py | 117 +-- litellm/types/llms/anthropic.py | 4 +- ...est_responses_bridge_sync_reasoning_sdk.py | 130 +++ .../test_fake_stream_reasoning_wire.py | 874 ++++++++++++++++++ .../test_streaming_handler.py | 66 -- ...al_pass_through_adapters_transformation.py | 2 +- .../test_streaming_iterator_combined_chunk.py | 96 +- .../llms/base_llm/test_base_model_iterator.py | 26 +- .../chat/test_bedrock_converse_handler.py | 228 +++++ tests/unit/llms/openai/test_openai.py | 75 +- .../chat/test_openai_like_chat_handler.py | 75 ++ .../test_snowflake_chat_transformation.py | 22 +- .../test_streaming_iterator_transformation.py | 97 +- 22 files changed, 1814 insertions(+), 306 deletions(-) create mode 100644 tests/integration/sdk/test_responses_bridge_sync_reasoning_sdk.py create mode 100644 tests/integration/streaming/test_fake_stream_reasoning_wire.py create mode 100644 tests/unit/llms/openai_like/chat/test_openai_like_chat_handler.py diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 95d062a190c..78fb82c98e7 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1128,32 +1128,24 @@ class CustomStreamWrapper: """ - if self.merge_reasoning_content_in_choices is True: - reasoning_content: Final = getattr(model_response.choices[0].delta, "reasoning_content", None) - if reasoning_content: - self._emitted_disqualifying_content = True - if self.sent_first_thinking_block is False: - # Ensure content is not None before concatenation - if model_response.choices[0].delta.content is None: - model_response.choices[0].delta.content = "" - model_response.choices[0].delta.content += "" + reasoning_content - self.sent_first_thinking_block = True - elif ( - self.sent_first_thinking_block is True - and hasattr(model_response.choices[0].delta, "reasoning_content") - and model_response.choices[0].delta.reasoning_content - ): - model_response.choices[0].delta.content = reasoning_content - elif ( - self.sent_first_thinking_block is True - and not self.sent_last_thinking_block - and model_response.choices[0].delta.content - ): - model_response.choices[0].delta.content = "" + (model_response.choices[0].delta.content or "") - self.sent_last_thinking_block = True + if self.merge_reasoning_content_in_choices is not True: + return + delta: Final = model_response.choices[0].delta + reasoning_content: Final = getattr(delta, "reasoning_content", None) + content: Final = delta.content or "" + if reasoning_content: + self._emitted_disqualifying_content = True + closes_thinking: Final = bool(content) and not self.sent_last_thinking_block + opening: Final = "" if self.sent_first_thinking_block else "" + delta.content = opening + reasoning_content + ("" + content if closes_thinking else "") + self.sent_first_thinking_block = True + self.sent_last_thinking_block = self.sent_last_thinking_block or closes_thinking + elif self.sent_first_thinking_block and not self.sent_last_thinking_block and content: + delta.content = "" + content + self.sent_last_thinking_block = True - if hasattr(model_response.choices[0].delta, "reasoning_content"): - del model_response.choices[0].delta.reasoning_content + if hasattr(delta, "reasoning_content"): + del delta.reasoning_content def _dispatch_provider_chunk( self, diff --git a/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py index 18db0f28338..52b836c789a 100644 --- a/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py @@ -24,6 +24,7 @@ from typing_extensions import assert_never from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.exceptions import ContextWindowExceededError, MidStreamFallbackError +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.hidden_params import set_hidden_params from litellm.llms.anthropic.pass_through.safeguards import StreamedSafeguardResults from litellm.llms.anthropic.pass_through.utils import anthropic_sse_frame @@ -39,6 +40,7 @@ from litellm.types.llms.anthropic import ( UsageDelta, UsageIteration, ) +from litellm.types.llms.openai import ChatCompletionRedactedThinkingBlock, ChatCompletionThinkingBlock from litellm.types.utils import AdapterCompletionStreamWrapper, ChatCompletionDeltaToolCall, Delta, Function if TYPE_CHECKING: @@ -50,6 +52,7 @@ if TYPE_CHECKING: _STREAMING_DELTA_TYPES: Final = frozenset(get_args(StreamingContentBlockDeltaType)) _SSE_EVENT: Final = TypeAdapter(Mapping[str, object]) +_THINKING_BLOCK: Final = TypeAdapter(ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock) class _UsageDeltaWithIterations(UsageDelta, total=False): @@ -106,6 +109,14 @@ def _delta_payload_field(delta_type: StreamingContentBlockDeltaType) -> str: assert_never(delta_type) +def _first_thinking_block_index(chunk: "ModelResponseStream") -> int | None: + delta: Final = _optional_attr(chunk.choices[0], "delta") if chunk.choices else None + blocks: Final = _optional_attr_sequence(delta, "thinking_blocks") + first_block: Final = as_str_mapping(blocks[0]) if blocks else None + index: Final = first_block.get("index") if first_block is not None else None + return index if isinstance(index, int) else None + + def _optional_str_attr(obj: object, name: str) -> str | None: value: Final = _optional_attr(obj, name) return value if isinstance(value, str) else None @@ -373,13 +384,44 @@ class _CombinedChunkSplitter: finish_delta.thinking_blocks = None return [content_chunk, finish_chunk] + @staticmethod + def _with_thinking_block( + chunk: "ModelResponseStream", + block: ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock, + *, + keep_role: bool, + ) -> "ModelResponseStream": + piece: Final = copy.deepcopy(chunk) + role: Final = getattr(piece.choices[0].delta, "role", None) if keep_role else None + thinking: Final = block.get("thinking") if block["type"] == "thinking" else None + piece.choices[0].delta = Delta(role=role, reasoning_content=thinking or None, thinking_blocks=[block]) + return piece + + @staticmethod + def _split_thinking_blocks(chunk: "ModelResponseStream") -> "tuple[ModelResponseStream, ...]": + """Anthropic streams one content block per thinking block, so a delta carrying several is split per block.""" + choices: Final = _optional_attr_sequence(chunk, "choices") + if len(choices) != 1: + return (chunk,) + delta: Final = _optional_attr(choices[0], "delta") + blocks: Final = _optional_attr_sequence(delta, "thinking_blocks") + if len(blocks) < 2 or _optional_attr(delta, "content") or _optional_attr(delta, "tool_calls"): + return (chunk,) + return tuple( + _CombinedChunkSplitter._with_thinking_block( + chunk, _THINKING_BLOCK.validate_python(block), keep_role=position == 0 + ) + for position, block in enumerate(blocks) + ) + def _expand(self, chunk: "ModelResponseStream") -> "tuple[ModelResponseStream, ...]": split_chunks: Final = tuple( split_chunk for combined_chunk in self._split(chunk) for split_chunk in self._split_by_payload_kind(combined_chunk) # comprehension-ok: flattens nested splits ) - return tuple(itertools.chain.from_iterable(self._route_tool_calls(split_chunk) for split_chunk in split_chunks)) + block_chunks: Final = tuple(itertools.chain.from_iterable(map(self._split_thinking_blocks, split_chunks))) + return tuple(itertools.chain.from_iterable(map(self._route_tool_calls, block_chunks))) def __iter__(self) -> "Iterator[ModelResponseStream]": return self @@ -458,7 +500,8 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): sent_first_chunk: bool = False sent_content_block_start: bool = False sent_content_block_finish: bool = False - current_content_block_type: Literal["text", "tool_use", "thinking"] = "text" + current_content_block_type: Literal["text", "tool_use", "thinking", "redacted_thinking"] = "text" + current_thinking_block_index: int | None = None sent_last_message: bool = False holding_chunk: ContentBlockDelta | None = None holding_stop_reason_chunk: MessageBlockDelta | None = None @@ -1254,13 +1297,14 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): Also gates re-emission after a block transition: when an upstream chunk both *triggers* a new content block (its type differs from the active block) and *carries* delta content, that content belongs to - the new block. The synthesized ``content_block_start`` only ever - carries an empty body — see - ``_translate_streaming_openai_chunk_to_anthropic_content_block``, - which returns an empty ``TextBlock``/``ToolUseBlock``/thinking block — - so the trigger chunk's delta must be re-queued or the first token of - the new block (the first non-empty text/thinking delta, or bundled - tool arguments) is silently dropped. + the new block. The synthesized ``content_block_start`` from + ``_translate_streaming_openai_chunk_to_anthropic_content_block`` + carries an empty text, tool-use or unsigned thinking body, and only a + signed thinking block's text rides on the start because its delta is + a ``signature_delta``. So the trigger chunk's delta must be re-queued + or the first token of the new block (the first non-empty text or + unsigned thinking delta, or bundled tool arguments) is silently + dropped. Delta types outside ``StreamingContentBlockDeltaType`` — the closed set the translate layer can produce — are treated as empty. The @@ -1352,11 +1396,24 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): original_name: Final = self.tool_name_mapping.get(truncated_name, truncated_name) tool_block["name"] = original_name + thinking_block_index: Final = _first_thinking_block_index(chunk) + previous_thinking_block_index: Final = self.current_thinking_block_index + if thinking_block_index is not None or block_type != self.current_content_block_type: + self.current_thinking_block_index = thinking_block_index + if block_type != self.current_content_block_type: self.current_content_block_type = block_type self.current_content_block_start = content_block_start return True + if ( + thinking_block_index is not None + and previous_thinking_block_index is not None + and thinking_block_index != previous_thinking_block_index + ): + self.current_content_block_start = content_block_start + return True + # For parallel tool calls, we'll necessarily have a new content block # if we get a function name since it signals a new tool call if block_type == "tool_use": diff --git a/litellm/llms/anthropic/pass_through/adapters/transformation.py b/litellm/llms/anthropic/pass_through/adapters/transformation.py index d1560c5b2ed..3aad0be1913 100644 --- a/litellm/llms/anthropic/pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/pass_through/adapters/transformation.py @@ -1627,7 +1627,7 @@ class LiteLLMAnthropicMessagesAdapter: def translate_streaming_openai_chunk_to_anthropic_content_block( self, choices: list[OpenAIStreamingChoice | StreamingChoices] ) -> tuple[ - Literal["text", "tool_use", "thinking"], + Literal["text", "tool_use", "thinking", "redacted_thinking"], "ContentBlockContentBlockDict", ]: from litellm._uuid import uuid @@ -1664,6 +1664,10 @@ class LiteLLMAnthropicMessagesAdapter: thinking_blocks = choice.delta.thinking_blocks or [] if len(thinking_blocks) > 0: thinking_block = thinking_blocks[0] + if thinking_block["type"] == "redacted_thinking": + return "redacted_thinking", ChatCompletionRedactedThinkingBlock( + type="redacted_thinking", data=thinking_block.get("data") or "" + ) if thinking_block["type"] == "thinking": thinking = thinking_block.get("thinking") or "" signature = thinking_block.get("signature") or "" @@ -1672,7 +1676,7 @@ class LiteLLMAnthropicMessagesAdapter: assert isinstance(signature, str) return "thinking", ChatCompletionThinkingBlock( - type="thinking", thinking=thinking, signature=signature + type="thinking", thinking=thinking if signature else "", signature=signature ) # OpenAI-compatible reasoning backends (e.g. vLLM/SGLang reasoning # parsers) populate ``reasoning_content`` without ``thinking_blocks``. diff --git a/litellm/llms/base_llm/base_model_iterator.py b/litellm/llms/base_llm/base_model_iterator.py index aede8aedc60..611051bb2dd 100644 --- a/litellm/llms/base_llm/base_model_iterator.py +++ b/litellm/llms/base_llm/base_model_iterator.py @@ -6,16 +6,38 @@ import litellm if TYPE_CHECKING: import httpx +from litellm.types.llms.openai import ChatCompletionRedactedThinkingBlock, ChatCompletionThinkingBlock from litellm.types.utils import ( Choices, Delta, GenericStreamingChunk, + Message, ModelResponse, ModelResponseStream, StreamingChoices, ) +def _indexed_thinking_block( + block: ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock, index: int +) -> ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock: + if block["type"] == "redacted_thinking": + return {**block, "index": index} + return {**block, "index": index} + + +def _streaming_delta(message: Message) -> Delta: + thinking_blocks: Final = message.thinking_blocks if hasattr(message, "thinking_blocks") else None + return Delta( + **message.model_dump(exclude={"thinking_blocks"}), + thinking_blocks=( + None + if thinking_blocks is None + else [_indexed_thinking_block(block, index) for index, block in enumerate(thinking_blocks)] + ), + ) + + def convert_model_response_to_streaming( model_response: ModelResponse, ) -> ModelResponseStream: @@ -40,9 +62,7 @@ def convert_model_response_to_streaming( streaming_choices.append( StreamingChoices( index=choice.index, - delta=Delta( - **cast(Choices, choice).message.model_dump(), - ), + delta=_streaming_delta(cast(Choices, choice).message), finish_reason=choice.finish_reason, ) ) diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 35d08d4b2c2..5a64e3fb78e 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -1,5 +1,5 @@ import json -from collections.abc import Iterator +from collections.abc import Iterator, Mapping from typing import Final import httpx @@ -10,6 +10,7 @@ from litellm.anthropic_beta_headers_manager import ( ) from litellm.litellm_core_utils.get_litellm_params import stored_control_options from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject +from litellm.llms.base_llm import base_model_iterator from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, @@ -21,7 +22,7 @@ from litellm.utils import CustomStreamWrapper from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token, pop_aws_auth_params, run_aws_signing from ..common_utils import BedrockError, error_response_text, get_all_bedrock_regions -from .invoke_handler import AWSEventStreamDecoder, MockResponseIterator, make_call +from .invoke_handler import AWSEventStreamDecoder, make_call def make_sync_call( @@ -36,7 +37,8 @@ def make_sync_call( fake_stream: bool = False, stream_chunk_size: int | None = None, timeout: float | httpx.Timeout | None = None, -) -> tuple[MockResponseIterator | Iterator[object], httpx.Headers]: + optional_params: Mapping[str, object] | None = None, +) -> tuple[base_model_iterator.MockResponseIterator | Iterator[object], httpx.Headers]: if client is None: client = get_httpx_client() # Create a new client if none provided @@ -64,14 +66,15 @@ def make_sync_call( model_response=litellm.ModelResponse(), stream=True, logging_obj=logging_obj, - optional_params={}, + optional_params=dict(optional_params or {}), api_key="", data=data, messages=messages, encoding=litellm.encoding, + json_mode=json_mode, ) - completion_stream: MockResponseIterator | Iterator[object] = MockResponseIterator( - model_response=model_response, json_mode=json_mode + completion_stream: base_model_iterator.MockResponseIterator | Iterator[object] = ( + base_model_iterator.MockResponseIterator(model_response=model_response) ) else: decoder: Final = AWSEventStreamDecoder(model=model, json_mode=json_mode) @@ -158,6 +161,7 @@ class BedrockConverseLLM(BaseAWSLLM): json_mode=json_mode, stream_chunk_size=stream_chunk_size, timeout=timeout, + optional_params=optional_params, ) streaming_response: Final = CustomStreamWrapper( completion_stream=completion_stream, @@ -290,7 +294,7 @@ class BedrockConverseLLM(BaseAWSLLM): stream_chunk_size: Final = stored_control_options(litellm_params).stream_chunk_size if stream is True else None unencoded_model_id: Final = optional_params.pop("model_id", None) fake_stream = optional_params.pop("fake_stream", False) - json_mode: Final = optional_params.get("json_mode", False) + json_mode: Final = optional_params.get("json_mode") is True if unencoded_model_id is not None: modelId = self.encode_model_id(model_id=unencoded_model_id) else: @@ -463,6 +467,7 @@ class BedrockConverseLLM(BaseAWSLLM): fake_stream=fake_stream, stream_chunk_size=stream_chunk_size, timeout=timeout, + optional_params=optional_params, ) streaming_response: Final = CustomStreamWrapper( completion_stream=completion_stream, diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 6ad400661be..7553dacd35d 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -1,7 +1,7 @@ import json import types from collections.abc import AsyncIterator, Iterator, Mapping -from typing import TYPE_CHECKING, Final, cast +from typing import TYPE_CHECKING, Final import httpx from pydantic import TypeAdapter, ValidationError @@ -20,6 +20,7 @@ from litellm.llms.anthropic.chat.handler import ( ModelResponseIterator as AnthropicModelResponseIterator, ) from litellm.llms.anthropic.chat.transformation import JsonModeRequest +from litellm.llms.base_llm import base_model_iterator from litellm.llms.base_llm.base_utils import ResponseFormatSchema from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -32,12 +33,8 @@ from litellm.types.llms.openai import ( ChatCompletionRedactedThinkingBlock, ChatCompletionThinkingBlock, ChatCompletionToolCallChunk, - ChatCompletionToolCallFunctionChunk, - ChatCompletionUsageBlock, ) from litellm.types.utils import ( - ChatCompletionMessageToolCall, - Choices, Delta, ModelResponse, ModelResponseStream, @@ -228,7 +225,10 @@ async def make_call( bedrock_invoke_provider: litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL | None = None, stream_chunk_size: int | None = None, timeout: float | httpx.Timeout | None = None, -) -> "tuple[MockResponseIterator | AsyncIterator[GChunk | ModelResponseStream | dict], httpx.Headers]": + optional_params: Mapping[str, object] | None = None, +) -> tuple[ + base_model_iterator.MockResponseIterator | AsyncIterator[GChunk | ModelResponseStream | dict], httpx.Headers +]: try: if client is None: client = get_async_httpx_client( @@ -264,15 +264,16 @@ async def make_call( model_response=litellm.ModelResponse(), stream=True, logging_obj=logging_obj, - optional_params={}, + optional_params=dict(optional_params or {}), api_key="", data=data, messages=messages, encoding=litellm.encoding, + json_mode=json_mode, ) - completion_stream: MockResponseIterator | AsyncIterator[GChunk | ModelResponseStream | dict] = ( - MockResponseIterator(model_response=model_response, json_mode=json_mode) - ) + completion_stream: ( + base_model_iterator.MockResponseIterator | AsyncIterator[GChunk | ModelResponseStream | dict] + ) = base_model_iterator.MockResponseIterator(model_response=model_response) else: decoder: Final = _stream_decoder( bedrock_invoke_provider, @@ -324,7 +325,7 @@ def make_sync_call( bedrock_invoke_provider: litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL | None = None, stream_chunk_size: int | None = None, timeout: float | httpx.Timeout | None = None, -) -> "tuple[MockResponseIterator | Iterator[GChunk | ModelResponseStream | dict], httpx.Headers]": +) -> tuple[base_model_iterator.MockResponseIterator | Iterator[GChunk | ModelResponseStream | dict], httpx.Headers]: try: if client is None: client = get_httpx_client( @@ -364,10 +365,11 @@ def make_sync_call( data=data, messages=messages, encoding=litellm.encoding, + json_mode=json_mode, ) - completion_stream: MockResponseIterator | Iterator[GChunk | ModelResponseStream | dict] = ( - MockResponseIterator(model_response=model_response, json_mode=json_mode) - ) + completion_stream: ( + base_model_iterator.MockResponseIterator | Iterator[GChunk | ModelResponseStream | dict] + ) = base_model_iterator.MockResponseIterator(model_response=model_response) else: decoder: Final = _stream_decoder( bedrock_invoke_provider, @@ -991,98 +993,3 @@ class AmazonOpenAICompatibleStreamDecoder(AWSEventStreamDecoder): def _chunk_parser(self, chunk_data: dict[str, object]) -> ModelResponseStream: return self.openai_model_response_iterator.chunk_parser(chunk=chunk_data) - - -class MockResponseIterator: # for returning ai21 streaming responses - def __init__(self, model_response, json_mode: bool | None = False): - self.model_response = model_response - self.json_mode = json_mode - self.is_done = False - - # Sync iterator - def __iter__(self): - return self - - def _handle_json_mode_chunk( - self, text: str, tool_calls: list[ChatCompletionToolCallChunk] | None - ) -> tuple[str, ChatCompletionToolCallChunk | None]: - """ - If JSON mode is enabled, convert the tool call to a message. - - Bedrock returns the JSON schema as part of the tool call - OpenAI returns the JSON schema as part of the content, this handles placing it in the content - - Args: - text: str - tool_use: Optional[ChatCompletionToolCallChunk] - Returns: - Tuple[str, Optional[ChatCompletionToolCallChunk]] - - text: The text to use in the content - tool_use: The ChatCompletionToolCallChunk to use in the chunk response - """ - tool_use: ChatCompletionToolCallChunk | None = None - if self.json_mode is True and tool_calls is not None: - message: Final = litellm.AnthropicConfig().convert_tool_response_to_message(tool_calls=tool_calls) - if message is not None: - text = message.content or "" - tool_use = None - elif tool_calls is not None and len(tool_calls) > 0: - tool_use = tool_calls[0] - return text, tool_use - - def _chunk_parser(self, chunk_data: ModelResponse) -> GChunk: - try: - chunk_usage: Final[Usage] = getattr(chunk_data, "usage") - text = chunk_data.choices[0].message.content or "" - tool_use = None - _model_response_tool_call: Final = cast( - list[ChatCompletionMessageToolCall] | None, - cast(Choices, chunk_data.choices[0]).message.tool_calls, - ) - if self.json_mode is True: - text, tool_use = self._handle_json_mode_chunk( - text=text, - tool_calls=chunk_data.choices[0].message.tool_calls, - ) - elif _model_response_tool_call is not None: - tool_use = ChatCompletionToolCallChunk( - id=_model_response_tool_call[0].id, - type="function", - function=ChatCompletionToolCallFunctionChunk( - name=_model_response_tool_call[0].function.name, - arguments=_model_response_tool_call[0].function.arguments, - ), - index=0, - ) - processed_chunk: Final = GChunk( - text=text, - tool_use=tool_use, - is_finished=True, - finish_reason=map_finish_reason(finish_reason=chunk_data.choices[0].finish_reason or ""), - usage=ChatCompletionUsageBlock( - prompt_tokens=chunk_usage.prompt_tokens, - completion_tokens=chunk_usage.completion_tokens, - total_tokens=chunk_usage.total_tokens, - ), - index=0, - ) - return processed_chunk - except Exception as e: - raise ValueError(f"Failed to decode chunk: {chunk_data}. Error: {e}") - - def __next__(self): - if self.is_done: - raise StopIteration - self.is_done = True - return self._chunk_parser(self.model_response) - - # Async iterator - def __aiter__(self): - return self - - async def __anext__(self): - if self.is_done: - raise StopAsyncIteration - self.is_done = True - return self._chunk_parser(self.model_response) diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 27346ee936e..79c69b5ceb4 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -30,9 +30,8 @@ from litellm.files.types import FileContentStreamingResult from litellm.litellm_core_utils.core_helpers import set_provider_response_headers_in_hidden_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_utils import speech_request_body, track_llm_api_timing -from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator +from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator, MockResponseIterator from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException -from litellm.llms.bedrock.chat.invoke_handler import MockResponseIterator from litellm.types.utils import ( EmbeddingResponse, ImageResponse, diff --git a/litellm/llms/openai_like/chat/handler.py b/litellm/llms/openai_like/chat/handler.py index 62a84ad3f15..2c74ef19e7d 100644 --- a/litellm/llms/openai_like/chat/handler.py +++ b/litellm/llms/openai_like/chat/handler.py @@ -14,7 +14,7 @@ from typing_extensions import ReadOnly import litellm from litellm import LlmProviders from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.llms.bedrock.chat.invoke_handler import MockResponseIterator +from litellm.llms.base_llm.base_model_iterator import MockResponseIterator from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.databricks.streaming_utils import ModelResponseIterator from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig diff --git a/litellm/llms/snowflake/chat/transformation.py b/litellm/llms/snowflake/chat/transformation.py index bf2a7d7da15..6a480904a3e 100644 --- a/litellm/llms/snowflake/chat/transformation.py +++ b/litellm/llms/snowflake/chat/transformation.py @@ -17,6 +17,7 @@ from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict import httpx from typing_extensions import ReadOnly +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.prompt_templates.factory import ( anthropic_process_openai_file_message, convert_to_anthropic_tool_result, @@ -184,10 +185,11 @@ def _signed_thinking_blocks(msg: object) -> list[dict[str, object]]: # mutable- blocks: Final = msg.get("thinking_blocks") if isinstance(msg, dict) else getattr(msg, "thinking_blocks", None) if not isinstance(blocks, list): return [] + mappings: Final = tuple(filter(None, map(as_str_mapping, blocks))) return [ - dict(block) - for block in blocks - if isinstance(block, Mapping) and (block.get("signature") or block.get("type") == "redacted_thinking") + {key: value for key, value in block.items() if key != "index"} + for block in mappings + if block.get("signature") or block.get("type") == "redacted_thinking" ] diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 418bdb3d6db..5c9e5cac5ab 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -100,6 +100,11 @@ def _delta_carries_output(delta: ChatCompletionDelta) -> bool: return any(fields) or _delta_has_signed_thinking_block(delta) +def _delta_carries_reasoning_and_answer(delta: ChatCompletionDelta) -> bool: + has_reasoning: Final = bool(getattr(delta, "reasoning_content", None)) or _delta_has_signed_thinking_block(delta) + return has_reasoning and bool(delta.content or delta.tool_calls) + + class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): """ Async iterator for processing streaming responses from the Responses API. @@ -566,6 +571,26 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._buffered_chunk = None return buffered + def _reasoning_part_first(self, chunk: ModelResponseStream) -> ModelResponseStream: + if not chunk.choices or not _delta_carries_reasoning_and_answer(chunk.choices[0].delta): + return chunk + first_choice: Final = chunk.choices[0] + delta: Final = first_choice.delta + answer_delta: Final = delta.model_copy(update={"reasoning_content": None, "thinking_blocks": None}) + self._buffered_chunk = chunk.model_copy( + update={"choices": [first_choice.model_copy(update={"delta": answer_delta}), *chunk.choices[1:]]} + ) + reasoning_delta: Final = ChatCompletionDelta( + reasoning_content=getattr(delta, "reasoning_content", None), + thinking_blocks=delta.thinking_blocks if hasattr(delta, "thinking_blocks") else None, + ) + return chunk.model_copy( + update={ + "provider_specific_fields": None, + "choices": [first_choice.model_copy(update={"delta": reasoning_delta, "finish_reason": None})], + } + ) + def _with_encoded_response_id(self, response: ResponsesAPIResponse) -> ResponsesAPIResponse: return ResponsesAPIRequestUtils.update_responses_api_response_id_with_model_id( responses_api_response=response, @@ -1052,6 +1077,41 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._queue_message_item_added_events() return + def _queue_reasoning_done_events_on_reasoning_end(self, chunk: ModelResponseStream) -> None: + if not self._reasoning_active or self._reasoning_done_emitted: + return + delta: Final = chunk.choices[0].delta if chunk.choices else None + if delta and getattr(delta, "reasoning_content", None): + self._accumulated_reasoning_content_parts.append(delta.reasoning_content) + if not self._is_reasoning_end(chunk): + return + reasoning_content: Final = "".join(self._accumulated_reasoning_content_parts) + self._cached_reasoning_item_id = ( + self._reasoning_item_id or self._cached_reasoning_item_id or mint_reasoning_item_id() + ) + reasoning_item_id: Final = self._cached_reasoning_item_id + self._sequence_number += 1 + text_done_event: Final = self.create_reasoning_summary_text_done_event( + reasoning_item_id=reasoning_item_id, + reasoning_content=reasoning_content, + sequence_number=self._sequence_number, + ) + self._sequence_number += 1 + part_done_event: Final = self.create_reasoning_summary_part_done_event( + reasoning_item_id=reasoning_item_id, + reasoning_content=reasoning_content, + sequence_number=self._sequence_number, + ) + self._sequence_number += 1 + reasoning_output_item_done_event: Final = self.create_reasoning_output_item_done_event( + reasoning_item_id=reasoning_item_id, + reasoning_content=reasoning_content, + sequence_number=self._sequence_number, + ) + self._pending_response_events.extend([text_done_event, part_done_event, reasoning_output_item_done_event]) + self._reasoning_done_emitted = True + self._reasoning_active = False + async def __anext__( self, ) -> ResponsesAPIStreamingResponse | ResponseCompletedEvent | BaseLiteLLMOpenAIResponseObject: @@ -1078,7 +1138,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): raise StopAsyncIteration chunk = await self.litellm_custom_stream_wrapper.__anext__() if chunk is not None: - chunk = cast(ModelResponseStream, chunk) + chunk = self._reasoning_part_first(cast(ModelResponseStream, chunk)) for src in ( getattr(chunk, "provider_specific_fields", None), getattr( @@ -1095,55 +1155,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self.collected_chat_completion_chunks.append( self._snapshot_chunk_for_stream_chunk_builder(chunk) ) - if self._reasoning_active and not self._reasoning_done_emitted: - # Incrementally accumulate reasoning content instead of - # calling stream_chunk_builder on every chunk (O(n²)) - delta = chunk.choices[0].delta if chunk.choices else None - if delta and hasattr(delta, "reasoning_content") and delta.reasoning_content: - self._accumulated_reasoning_content_parts.append(delta.reasoning_content) - if self._is_reasoning_end(chunk): - reasoning_content = "".join(self._accumulated_reasoning_content_parts) - - # Ensure we have a valid reasoning_item_id - self._cached_reasoning_item_id = ( - self._reasoning_item_id - or self._cached_reasoning_item_id - or mint_reasoning_item_id() - ) - reasoning_item_id = self._cached_reasoning_item_id - - # Create text.done event first with its own sequence number - self._sequence_number += 1 - text_done_event = self.create_reasoning_summary_text_done_event( - reasoning_item_id=reasoning_item_id, - reasoning_content=reasoning_content, - sequence_number=self._sequence_number, - ) - - # Create part.done event second with its own sequence number - self._sequence_number += 1 - part_done_event = self.create_reasoning_summary_part_done_event( - reasoning_item_id=reasoning_item_id, - reasoning_content=reasoning_content, - sequence_number=self._sequence_number, - ) - - self._sequence_number += 1 - reasoning_output_item_done_event = self.create_reasoning_output_item_done_event( - reasoning_item_id=reasoning_item_id, - reasoning_content=reasoning_content, - sequence_number=self._sequence_number, - ) - self._pending_response_events.extend( - [ - text_done_event, - part_done_event, - reasoning_output_item_done_event, - ] - ) - self._reasoning_done_emitted = True - self._reasoning_active = False - + self._queue_reasoning_done_events_on_reasoning_end(chunk) response_api_chunk = self._transform_chat_completion_chunk_to_response_api_chunk(chunk) if response_api_chunk: self._pending_response_events.append(response_api_chunk) @@ -1182,11 +1194,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): try: buffered_chunk = self._take_buffered_chunk() if buffered_chunk is not None: - chunk = buffered_chunk + chunk = self._reasoning_part_first(buffered_chunk) elif self._upstream_exhausted: raise StopIteration else: - chunk = self.litellm_custom_stream_wrapper.__next__() + chunk = self._reasoning_part_first(self.litellm_custom_stream_wrapper.__next__()) for src in ( getattr(chunk, "provider_specific_fields", None), getattr( @@ -1206,6 +1218,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self.collected_chat_completion_chunks.append( self._snapshot_chunk_for_stream_chunk_builder(cast(ModelResponseStream, chunk)) ) + self._queue_reasoning_done_events_on_reasoning_end(chunk) response_api_chunk = self._transform_chat_completion_chunk_to_response_api_chunk(chunk) if response_api_chunk: self._pending_response_events.append(response_api_chunk) diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index c3d257117ed..6aeead7679f 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -536,7 +536,9 @@ class ContentBlockStartText(TypedDict): content_block: TextBlock -ContentBlockContentBlockDict = ToolUseBlock | TextBlock | ChatCompletionThinkingBlock +ContentBlockContentBlockDict = ( + ToolUseBlock | TextBlock | ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock +) ContentBlockStart = ContentBlockStartToolUse | ContentBlockStartText diff --git a/tests/integration/sdk/test_responses_bridge_sync_reasoning_sdk.py b/tests/integration/sdk/test_responses_bridge_sync_reasoning_sdk.py new file mode 100644 index 00000000000..99f5fa941f7 --- /dev/null +++ b/tests/integration/sdk/test_responses_bridge_sync_reasoning_sdk.py @@ -0,0 +1,130 @@ +import json +import uuid +from collections.abc import Iterable, Mapping, Sequence +from typing import Final + +import pytest +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter + +import litellm +from litellm.responses.litellm_completion_transformation.streaming_iterator import LiteLLMCompletionStreamingIterator + +_BACKEND: Final = "qwen3-sync-bridge-audit" +_API_KEY: Final = "synthetic-hosted-vllm-key" +_ANSWER: Final = "It is sunny in Paris." +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_ITEM_EVENTS: Final = frozenset({"response.output_item.added", "response.output_item.done"}) + + +class _Frozen(BaseModel): + model_config = ConfigDict(frozen=True) + + +class _Item(_Frozen): + type: str + + +class _Event(_Frozen): + type: str + item: _Item | None = None + output_index: int | None = None + delta: JsonValue = None + + +def _frame(identity: str, delta: Mapping[str, JsonValue], finish_reason: str | None = None) -> bytes: + usage: Final = ( + {"usage": {"prompt_tokens": 30, "completion_tokens": 12, "total_tokens": 42}} if finish_reason else {} + ) + body: Final = { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": _BACKEND, + "choices": [{"index": 0, "delta": dict(delta), "finish_reason": finish_reason}], + **usage, + } + return b"data: " + json.dumps(body).encode() + b"\n\n" + + +def _sse(deltas: Sequence[Mapping[str, JsonValue]]) -> Reply: + identity: Final = f"chatcmpl-sync-bridge-{uuid.uuid4().hex}" + return Reply( + content_type="text/event-stream", + chunks=( + _frame(identity, {"role": "assistant", "content": ""}), + *(_frame(identity, delta) for delta in deltas), + _frame(identity, {}, "stop"), + b"data: [DONE]\n\n", + ), + ) + + +def _streamed_chat_posts(wire: Wire) -> list[JsonValue]: + posted: Final = [request for request in wire.drain() if request.method == "POST"] + assert [request.target for request in posted] == ["/v1/chat/completions"] + return [_JSON_OBJECT.validate_json(request.body).get("stream") for request in posted] + + +def _events(stream: Iterable[object]) -> tuple[_Event, ...]: + return tuple(_Event.model_validate(_JSON_OBJECT.validate_python(_dumped(event))) for event in stream) + + +def _dumped(event: object) -> object: + assert isinstance(event, BaseModel), type(event) + return event.model_dump(mode="json") + + +def _delta_text(events: Sequence[_Event], event_type: str) -> str: + return "".join(str(event.delta) for event in events if event.type == event_type) + + +@pytest.mark.parametrize( + "deltas", + [ + pytest.param( + ({"reasoning_content": "REASONING"}, {"content": _ANSWER}), + id="reasoning-then-answer-chunks", + ), + pytest.param(({"reasoning_content": "REASONING", "content": _ANSWER},), id="reasoning-and-answer-in-one-chunk"), + ], +) +def test_sync_responses_stream_closes_the_reasoning_item_before_the_message_opens( + deltas: tuple[dict[str, JsonValue], ...], +) -> None: + marker: Final = uuid.uuid4().hex + reasoning: Final = f"Sync bridge reasoning {marker}." + scripted: Final = tuple( + {key: reasoning if key == "reasoning_content" else value for key, value in delta.items()} for delta in deltas + ) + + def respond(request: Request) -> Reply: + return Reply(body=b'{"object":"list","data":[]}') if request.method == "GET" else _sse(scripted) + + with wire_server(respond) as wire: + stream: Final = litellm.responses( + model=f"hosted_vllm/{_BACKEND}", + input=f"Weather in Paris? {marker}", + stream=True, + api_base=wire.url + "/v1", + api_key=_API_KEY, + use_chat_completions_api=True, + num_retries=0, + ) + assert isinstance(stream, LiteLLMCompletionStreamingIterator), type(stream) + events: Final = _events(stream) + + assert _streamed_chat_posts(wire) == [True] + lifecycle: Final = [ + (event.type, event.item.type, event.output_index) + for event in events + if event.type in _ITEM_EVENTS and event.item is not None + ] + assert lifecycle == [ + ("response.output_item.added", "reasoning", 0), + ("response.output_item.done", "reasoning", 0), + ("response.output_item.added", "message", 1), + ("response.output_item.done", "message", 1), + ], [event.type for event in events] + assert _delta_text(events, "response.reasoning_summary_text.delta") == reasoning + assert _delta_text(events, "response.output_text.delta") == _ANSWER diff --git a/tests/integration/streaming/test_fake_stream_reasoning_wire.py b/tests/integration/streaming/test_fake_stream_reasoning_wire.py new file mode 100644 index 00000000000..e46208365dd --- /dev/null +++ b/tests/integration/streaming/test_fake_stream_reasoning_wire.py @@ -0,0 +1,874 @@ +import base64 +import json +import math +import os +import uuid +from collections.abc import Callable, Generator, Mapping, Sequence +from contextlib import contextmanager +from dataclasses import dataclass +from importlib.resources import files +from itertools import chain +from typing import Final +from urllib.parse import unquote + +import anthropic +import openai +import pytest +from integration._support.client import Gateway, Scenario, eventually +from integration._support.database import read_rows +from integration._support.upstream import ( + _aws_event_frame, # pyright: ignore[reportPrivateUsage] # the scripted Bedrock peer's only event-stream encoder +) +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter +from redis import Redis + +_TOKEN: Final = "synthetic-bedrock-bearer" +_NATIVE_REDACTED: Final = base64.b64encode(b"\x12\x1esynthetic redacted reasoning payload").decode() +_CONVERSE_REDACTED: Final = base64.b64encode(_NATIVE_REDACTED.encode()).decode() +_SIGNATURE: Final = "EqoBCkgIBxABGAIiQsyntheticSignature" +_ANSWER: Final[dict[str, JsonValue]] = {"city": "Tokyo", "population_millions": 14.0} +_THINKING_TOKENS: Final = 87 +_EVENT_STREAM: Final = "application/vnd.amazon.eventstream" +_NO_CACHE: Final[dict[str, JsonValue]] = {"no-cache": True} +_THINKING: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": 1024} +_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"city": {"type": "string"}, "population_millions": {"type": "number"}}, + "required": ["city", "population_millions"], + "additionalProperties": False, +} +_RESPONSE_FORMAT: Final[dict[str, JsonValue]] = { + "type": "json_schema", + "json_schema": {"name": "city_fact", "schema": _SCHEMA}, +} +_CONVERSE_USAGE: Final[dict[str, JsonValue]] = { + "inputTokens": 40, + "outputTokens": 120, + "cacheReadInputTokens": 50, + "totalTokens": 210, +} +_LLAMA_MODEL_ID: Final = "us.meta.llama3-3-70b-instruct-v1:0" +_VLLM_BACKEND: Final = "qwen3-fake-stream-audit" +_VLLM_KEY: Final = "synthetic-hosted-vllm-key" +_AZURE_DEPLOYMENT: Final = "reasoning-deployment" +_AZURE_API_VERSION: Final = "2025-03-01-preview" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_DISCOVERY_PROBE: Final = ("GET", "/v1/models") + + +def _tool(name: str) -> dict[str, JsonValue]: + return { + "type": "function", + "function": { + "name": name, + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + }, + } + + +class _Frozen(BaseModel): + model_config = ConfigDict(frozen=True) + + +class _CostEntry(_Frozen): + litellm_provider: str | None = None + supports_native_structured_output: bool | None = None + supports_reasoning: bool | None = None + + +class _ToolFunction(_Frozen): + name: str | None = None + arguments: str | None = None + + +class _ToolCall(_Frozen): + index: int = 0 + function: _ToolFunction + + +class _PromptDetails(_Frozen): + cached_tokens: int | None = None + + +class _CompletionDetails(_Frozen): + reasoning_tokens: int | None = None + + +class _Usage(_Frozen): + prompt_tokens: int + completion_tokens: int + total_tokens: int + prompt_tokens_details: _PromptDetails | None = None + completion_tokens_details: _CompletionDetails | None = None + + +class _Message(_Frozen): + content: str | None = None + reasoning_content: str | None = None + thinking_blocks: tuple[dict[str, JsonValue], ...] | None = None + tool_calls: tuple[_ToolCall, ...] | None = None + + +class _Choice(_Frozen): + message: _Message + finish_reason: str | None = None + + +class _Completion(_Frozen): + choices: tuple[_Choice, ...] + usage: _Usage + + +class _StreamChoice(_Frozen): + delta: _Message + finish_reason: str | None = None + + +class _Chunk(_Frozen): + choices: tuple[_StreamChoice, ...] = () + usage: _Usage | None = None + + +class _ConverseMessage(_Frozen): + role: str + content: tuple[dict[str, JsonValue], ...] + + +class _ConverseRequest(_Frozen): + messages: tuple[_ConverseMessage, ...] + + +@dataclass(frozen=True, slots=True) +class _Answer: + content: str + reasoning: str + thinking_blocks: tuple[dict[str, JsonValue], ...] + tool_calls: tuple[tuple[str, JsonValue], ...] + finish_reason: str | None + usage: _Usage | None + + +_COST_MAP: Final = TypeAdapter(dict[str, _CostEntry]) + + +def _fakes_structured_output_stream(model: str, entry: _CostEntry) -> bool: + return ( + model.startswith("us.anthropic.claude") + and entry.litellm_provider == "bedrock_converse" + and entry.supports_native_structured_output is False + and entry.supports_reasoning is True + ) + + +def _claude_without_native_structured_output() -> str: + entries: Final = _COST_MAP.validate_json( + (files("litellm") / "model_prices_and_context_window_backup.json").read_bytes() + ) + return next(model for model, entry in entries.items() if _fakes_structured_output_stream(model, entry)) + + +def _converse_reply(content: Sequence[JsonValue], stop_reason: str) -> Reply: + return Reply( + body=json.dumps( + { + "output": {"message": {"role": "assistant", "content": list(content)}}, + "stopReason": stop_reason, + "usage": _CONVERSE_USAGE, + "additionalModelResponseFields": { + "usage": {"output_tokens_details": {"thinking_tokens": _THINKING_TOKENS}} + }, + "metrics": {"latencyMs": 1}, + } + ).encode() + ) + + +def _reasoned_json_answer(reasoning: str) -> list[JsonValue]: + return [ + {"reasoningContent": {"reasoningText": {"text": reasoning, "signature": _SIGNATURE}}}, + {"reasoningContent": {"redactedContent": _CONVERSE_REDACTED}}, + {"toolUse": {"toolUseId": "tooluse_json_1", "name": "json_tool_call", "input": _ANSWER}}, + ] + + +def _is_replay(request: Request) -> bool: + return any(message.role == "assistant" for message in _ConverseRequest.model_validate_json(request.body).messages) + + +def _converse_peer(reasoning: str) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + assert unquote(request.target).endswith("/converse"), request.target + if _is_replay(request): + return _converse_reply([{"text": "replayed"}], "end_turn") + return _converse_reply(_reasoned_json_answer(reasoning), "tool_use") + + return respond + + +def _converse_targets(wire: Wire) -> list[str]: + return [unquote(request.target).rsplit("/", 1)[-1] for request in wire.drain()] + + +def _claude_model(scenario: Scenario, wire: Wire) -> str: + return scenario.model( + model=f"bedrock/converse/{_claude_without_native_structured_output()}", + api_key=_TOKEN, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint=wire.url, + ) + + +def _openai_client(gateway: Gateway) -> openai.OpenAI: + return openai.OpenAI(base_url=str(gateway.client.base_url) + "/v1", api_key=gateway.key, max_retries=0) + + +def _parsed_arguments(arguments: str | None) -> JsonValue: + return _JSON_OBJECT.validate_json(arguments or "{}") + + +def _completion_answer(completion: _Completion) -> _Answer: + choice: Final = completion.choices[0] + return _Answer( + content=choice.message.content or "", + reasoning=choice.message.reasoning_content or "", + thinking_blocks=choice.message.thinking_blocks or (), + tool_calls=tuple( + (call.function.name or "", _parsed_arguments(call.function.arguments)) + for call in choice.message.tool_calls or () + ), + finish_reason=choice.finish_reason, + usage=completion.usage, + ) + + +def _joined_call(calls: Sequence[_ToolCall], index: int) -> tuple[str, JsonValue]: + pieces: Final = tuple(call.function for call in calls if call.index == index) + return ( + "".join(piece.name or "" for piece in pieces), + _parsed_arguments("".join(piece.arguments or "" for piece in pieces)), + ) + + +def _streamed_answer(chunks: Sequence[_Chunk]) -> _Answer: + choices: Final = tuple(chain.from_iterable(chunk.choices for chunk in chunks)) + deltas: Final = tuple(choice.delta for choice in choices) + calls: Final = tuple(chain.from_iterable(delta.tool_calls or () for delta in deltas)) + finishes: Final = tuple(choice.finish_reason for choice in choices if choice.finish_reason) + usages: Final = tuple(chunk.usage for chunk in chunks if chunk.usage is not None) + return _Answer( + content="".join(delta.content or "" for delta in deltas), + reasoning="".join(delta.reasoning_content or "" for delta in deltas), + thinking_blocks=tuple(chain.from_iterable(delta.thinking_blocks or () for delta in deltas)), + tool_calls=tuple(_joined_call(calls, index) for index in sorted({call.index for call in calls})), + finish_reason=finishes[-1] if finishes else None, + usage=usages[-1] if usages else None, + ) + + +def _chat(gateway: Gateway, body: Mapping[str, JsonValue]) -> _Answer: + completion: Final = _openai_client(gateway).chat.completions.create( + model=str(body["model"]), + messages=[{"role": "user", "content": str(body["question"])}], + extra_body={key: value for key, value in body.items() if key not in ("model", "question")}, + ) + return _completion_answer(_Completion.model_validate(completion.model_dump(mode="json"))) + + +def _chat_stream(gateway: Gateway, body: Mapping[str, JsonValue]) -> _Answer: + stream: Final = _openai_client(gateway).chat.completions.create( + model=str(body["model"]), + messages=[{"role": "user", "content": str(body["question"])}], + stream=True, + stream_options={"include_usage": True}, + extra_body={key: value for key, value in body.items() if key not in ("model", "question")}, + ) + return _streamed_answer(tuple(_Chunk.model_validate(chunk.model_dump(mode="json")) for chunk in stream)) + + +def _unindexed(blocks: Sequence[dict[str, JsonValue]]) -> tuple[dict[str, JsonValue], ...]: + return tuple({key: value for key, value in block.items() if key != "index"} for block in blocks) + + +def _json_question(model: str, marker: str, cache: Mapping[str, JsonValue] = _NO_CACHE) -> dict[str, JsonValue]: + return { + "model": model, + "question": f"How many people live in Tokyo, in millions? {marker}", + "response_format": _RESPONSE_FORMAT, + "thinking": _THINKING, + "max_tokens": 2048, + "cache": dict(cache), + } + + +def test_fake_streamed_json_schema_chat_answer_matches_the_non_streamed_answer(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + reasoning: Final = f"Tokyo proper has about 14 million people, marker {marker}." + with wire_server(_converse_peer(reasoning)) as wire, gateway.scenario() as scenario: + model: Final = _claude_model(scenario, wire) + question: Final = _json_question(model, marker) + + answer: Final = _chat(gateway, question) + streamed: Final = _chat_stream(gateway, question) + + assert _converse_targets(wire) == ["converse", "converse"] + assert (answer.reasoning, answer.finish_reason, answer.tool_calls) == (reasoning, "stop", ()) + assert _JSON_OBJECT.validate_json(answer.content) == _ANSWER + assert answer.thinking_blocks == ( + {"type": "thinking", "thinking": reasoning, "signature": _SIGNATURE}, + {"type": "redacted_thinking", "data": _NATIVE_REDACTED}, + ) + assert (streamed.reasoning, streamed.finish_reason, streamed.tool_calls) == (reasoning, "stop", ()) + assert _JSON_OBJECT.validate_json(streamed.content) == _ANSWER + assert _unindexed(streamed.thinking_blocks) == answer.thinking_blocks + assert streamed.usage == answer.usage + assert answer.usage is not None and answer.usage.completion_tokens_details is not None + assert answer.usage.completion_tokens_details.reasoning_tokens == _THINKING_TOKENS + + replayed: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "thinking": _THINKING, + "max_tokens": 2048, + "cache": _NO_CACHE, + "messages": [ + {"role": "user", "content": str(question["question"])}, + { + "role": "assistant", + "content": streamed.content, + "thinking_blocks": list(streamed.thinking_blocks), + }, + {"role": "user", "content": "And Osaka?"}, + ], + }, + ) + assert replayed.status_code == 200, replayed.text + sent: Final = _ConverseRequest.model_validate_json(wire.drain()[-1].body) + assert next(message for message in sent.messages if message.role == "assistant").content == ( + {"reasoningContent": {"reasoningText": {"text": reasoning, "signature": _SIGNATURE}}}, + {"reasoningContent": {"redactedContent": _CONVERSE_REDACTED}}, + {"text": streamed.content}, + ) + + +def _spend_rows(model: str, count: int) -> list[dict[str, JsonValue]]: + return eventually( + lambda: read_rows( + 'SELECT spend, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE model_group=%s', (model,) + ), + lambda found: len(found) >= count, + seconds=70, + ) + + +def test_fake_streamed_cached_prompt_is_billed_like_the_non_streamed_answer(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_converse_peer(f"Billing check {marker}.")) as wire, gateway.scenario() as scenario: + model: Final = _claude_model(scenario, wire) + question: Final = _json_question(model, marker) + _chat(gateway, question) + _chat_stream(gateway, question) + + rows: Final = _spend_rows(model, 2) + assert _converse_targets(wire) == ["converse", "converse"] + assert len(rows) == 2, rows + first, second = rows + assert (first["prompt_tokens"], first["completion_tokens"]) == ( + second["prompt_tokens"], + second["completion_tokens"], + ) + assert float(str(first["spend"])) > 0 + assert math.isclose(float(str(second["spend"])), float(str(first["spend"])), rel_tol=1e-9), rows + + +def _cache_holds(needle: str) -> bool: + with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as store: + return any(needle.encode() in _bytes(store.get(key)) for key in store.scan_iter(count=1000, _type="STRING")) + + +def _bytes(value: object) -> bytes: + return value if isinstance(value, bytes) else b"" + + +def test_cached_fake_stream_replays_the_answer_without_reaching_bedrock(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + reasoning: Final = f"Cached reasoning {marker}." + with wire_server(_converse_peer(reasoning)) as wire, gateway.scenario() as scenario: + model: Final = _claude_model(scenario, wire) + question: Final = _json_question(model, marker, cache={}) + + first: Final = _chat_stream(gateway, question) + assert first.reasoning == reasoning + eventually(lambda: _cache_holds(marker), bool, seconds=30) + assert _converse_targets(wire) == ["converse"] + + replay: Final = _chat_stream(gateway, question) + + assert _converse_targets(wire) == [] + assert (replay.content, replay.finish_reason, replay.tool_calls) == (first.content, "stop", ()) + assert _unindexed(replay.thinking_blocks) in ((), _unindexed(first.thinking_blocks)) + + +def _responses_reasoning_and_answer( + events: Sequence[dict[str, JsonValue]], +) -> tuple[list[tuple[str, str, JsonValue]], str, str]: + lifecycle: Final = [ + (str(event["type"]), str(_JSON_OBJECT.validate_python(event["item"])["type"]), event["output_index"]) + for event in events + if event["type"] in ("response.output_item.added", "response.output_item.done") + ] + summary: Final = "".join( + str(event["delta"]) for event in events if event["type"] == "response.reasoning_summary_text.delta" + ) + text: Final = "".join(str(event["delta"]) for event in events if event["type"] == "response.output_text.delta") + return lifecycle, summary, text + + +_REASONING_THEN_MESSAGE: Final = [ + ("response.output_item.added", "reasoning", 0), + ("response.output_item.done", "reasoning", 0), + ("response.output_item.added", "message", 1), + ("response.output_item.done", "message", 1), +] + + +def _raw_responses_events(gateway: Gateway, body: Mapping[str, JsonValue]) -> list[dict[str, JsonValue]]: + with gateway.client.stream( + "POST", + "/v1/responses", + json={**body, "stream": True, "store": False, "cache": _NO_CACHE}, + headers={"Authorization": f"Bearer {gateway.key}"}, + ) as response: + assert response.status_code == 200, response.read() + return [ + _JSON_OBJECT.validate_json(line.removeprefix("data: ")) + for line in response.iter_lines() + if line.startswith("data: {") + ] + + +def _completed_usage(events: Sequence[dict[str, JsonValue]]) -> dict[str, JsonValue]: + completed: Final = [event for event in events if event["type"] == "response.completed"] + assert len(completed) == 1, [event["type"] for event in events] + return _JSON_OBJECT.validate_python(_JSON_OBJECT.validate_python(completed[0]["response"])["usage"]) + + +def test_fake_streamed_json_schema_responses_stream_closes_the_reasoning_before_the_answer(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + reasoning: Final = f"Responses reasoning {marker}." + with wire_server(_converse_peer(reasoning)) as wire, gateway.scenario() as scenario: + model: Final = _claude_model(scenario, wire) + events: Final = _raw_responses_events( + gateway, + { + "model": model, + "input": f"How many people live in Tokyo, in millions? {marker}", + "reasoning": {"effort": "high"}, + "max_output_tokens": 2048, + "text": {"format": {"type": "json_schema", "name": "city_fact", "schema": _SCHEMA}}, + }, + ) + + lifecycle, summary, text = _responses_reasoning_and_answer(events) + assert _converse_targets(wire) == ["converse"] + assert lifecycle == _REASONING_THEN_MESSAGE, [event["type"] for event in events] + assert summary == reasoning + assert _JSON_OBJECT.validate_json(text) == _ANSWER + usage: Final = _completed_usage(events) + assert _JSON_OBJECT.validate_python(usage["input_tokens_details"])["cached_tokens"] == 50 + assert _JSON_OBJECT.validate_python(usage["output_tokens_details"])["reasoning_tokens"] == _THINKING_TOKENS + + +def _anthropic_client(gateway: Gateway) -> anthropic.Anthropic: + return anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) + + +def _block_fields(block: anthropic.types.ContentBlock) -> tuple[str, str, str]: + match block.type: + case "thinking": + return (block.type, block.thinking, block.signature) + case "redacted_thinking": + return (block.type, block.data, "") + case "text": + return (block.type, block.text, "") + case _: + return (block.type, "", "") + + +def _message_blocks(message: anthropic.types.Message) -> list[tuple[str, str, str]]: + return [_block_fields(block) for block in message.content] + + +def test_fake_streamed_json_schema_messages_stream_opens_one_block_per_thinking_block(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + reasoning: Final = f"Messages reasoning {marker}." + with wire_server(_converse_peer(reasoning)) as wire, gateway.scenario() as scenario: + model: Final = _claude_model(scenario, wire) + with _anthropic_client(gateway).messages.stream( + model=model, + max_tokens=2048, + thinking={"type": "enabled", "budget_tokens": 1024}, + messages=[{"role": "user", "content": f"How many people live in Tokyo, in millions? {marker}"}], + extra_body={"output_format": {"type": "json_schema", "schema": _SCHEMA}, "cache": _NO_CACHE}, + ) as stream: + final: Final = stream.get_final_message() + + assert _converse_targets(wire) == ["converse"] + blocks: Final = _message_blocks(final) + assert blocks[:2] == [ + ("thinking", reasoning, _SIGNATURE), + ("redacted_thinking", _NATIVE_REDACTED, ""), + ], blocks + assert [block[0] for block in blocks[2:]] == ["text"], blocks + assert _JSON_OBJECT.validate_json(blocks[2][1]) == _ANSWER + assert final.stop_reason == "end_turn" + + +def _frame(event_type: str, payload: Mapping[str, JsonValue]) -> bytes: + return _aws_event_frame(event_type, payload, "sc", "u") + + +def _converse_stream_with_thinking(reasoning_parts: Sequence[str], answer: str) -> Reply: + return Reply( + content_type=_EVENT_STREAM, + body=b"".join( + ( + _frame("messageStart", {"role": "assistant"}), + *( + _frame( + "contentBlockDelta", + {"delta": {"reasoningContent": {"text": part}}, "contentBlockIndex": 0}, + ) + for part in reasoning_parts + ), + _frame( + "contentBlockDelta", + {"delta": {"reasoningContent": {"signature": _SIGNATURE}}, "contentBlockIndex": 0}, + ), + _frame("contentBlockStop", {"contentBlockIndex": 0}), + _frame("contentBlockDelta", {"delta": {"text": answer}, "contentBlockIndex": 1}), + _frame("contentBlockStop", {"contentBlockIndex": 1}), + _frame("messageStop", {"stopReason": "end_turn"}), + _frame("metadata", {"usage": {"inputTokens": 11, "outputTokens": 9, "totalTokens": 20}}), + ) + ), + ) + + +def test_native_converse_stream_thinking_reaches_the_anthropic_sdk_once_with_its_signature(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + parts: Final = ("Native thinking ", f"marker {marker}.") + answer: Final = "It is sunny in Paris." + + def respond(request: Request) -> Reply: + assert unquote(request.target).endswith("/converse-stream"), request.target + return _converse_stream_with_thinking(parts, answer) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = _claude_model(scenario, wire) + with _anthropic_client(gateway).messages.stream( + model=model, + max_tokens=2048, + thinking={"type": "enabled", "budget_tokens": 1024}, + messages=[{"role": "user", "content": f"Weather in Paris? {marker}"}], + extra_body={"cache": _NO_CACHE}, + ) as stream: + final: Final = stream.get_final_message() + + assert _converse_targets(wire) == ["converse-stream"] + assert _message_blocks(final) == [("thinking", "".join(parts), _SIGNATURE), ("text", answer, "")] + + +def _llama_peer(content: list[JsonValue], stop_reason: str) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + assert unquote(request.target).endswith("/converse"), request.target + return Reply( + body=json.dumps( + { + "output": {"message": {"role": "assistant", "content": content}}, + "stopReason": stop_reason, + "usage": {"inputTokens": 175, "outputTokens": 23, "totalTokens": 198}, + "metrics": {"latencyMs": 1}, + } + ).encode() + ) + + return respond + + +_PARIS: Final[dict[str, JsonValue]] = {"city": "Paris"} + + +@pytest.mark.parametrize( + "content, stop_reason, expected_calls", + [ + pytest.param( + [ + {"toolUse": {"toolUseId": "tooluse_weather", "name": "get_weather", "input": _PARIS}}, + {"toolUse": {"toolUseId": "tooluse_time", "name": "get_time", "input": _PARIS}}, + ], + "tool_use", + (("get_weather", _PARIS), ("get_time", _PARIS)), + id="two-tool-use-blocks", + ), + pytest.param( + [{"text": json.dumps({"type": "function", "name": "get_weather", "parameters": _PARIS})}], + "end_turn", + (("get_weather", _PARIS),), + id="tool-call-written-as-text", + ), + ], +) +def test_fake_streamed_llama_tool_calls_match_the_non_streamed_answer( + gateway: Gateway, + content: list[JsonValue], + stop_reason: str, + expected_calls: tuple[tuple[str, JsonValue], ...], +) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_llama_peer(content, stop_reason)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"bedrock/converse/{_LLAMA_MODEL_ID}", + api_key=_TOKEN, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint=wire.url, + ) + question: Final[dict[str, JsonValue]] = { + "model": model, + "question": f"Weather and time in Paris? {marker}", + "tools": [_tool("get_weather"), _tool("get_time")], + "cache": _NO_CACHE, + } + + answer: Final = _chat(gateway, question) + streamed: Final = _chat_stream(gateway, question) + + assert _converse_targets(wire) == ["converse", "converse"] + assert (answer.tool_calls, answer.finish_reason) == (expected_calls, "tool_calls") + assert (streamed.tool_calls, streamed.finish_reason, streamed.content) == ( + answer.tool_calls, + "tool_calls", + answer.content, + ) + + +def _azure_o_series_reply() -> bytes: + return json.dumps( + { + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion", + "created": 1, + "model": "o4-mini", + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_weather", + "type": "function", + "function": {"name": "get_weather", "arguments": json.dumps(_PARIS)}, + }, + { + "id": "call_time", + "type": "function", + "function": {"name": "get_time", "arguments": json.dumps(_PARIS)}, + }, + ], + }, + } + ], + "usage": { + "prompt_tokens": 85, + "completion_tokens": 1955, + "total_tokens": 2040, + "prompt_tokens_details": {"cached_tokens": 64}, + "completion_tokens_details": {"reasoning_tokens": 1920}, + }, + } + ).encode() + + +def test_azure_o_series_fake_stream_keeps_every_tool_call_and_token_detail(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(lambda _: Reply(body=_azure_o_series_reply())) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"azure/o_series/{_AZURE_DEPLOYMENT}", + api_base=wire.url, + api_key="synthetic-azure-key", + api_version=_AZURE_API_VERSION, + ) + streamed: Final = _chat_stream( + gateway, + { + "model": model, + "question": f"Weather and time in Paris? {marker}", + "tools": [_tool("get_weather"), _tool("get_time")], + "cache": _NO_CACHE, + }, + ) + + received: Final = wire.drain() + assert [unquote(request.target).split("?")[0] for request in received] == [ + f"/openai/deployments/{_AZURE_DEPLOYMENT}/chat/completions" + ] + assert _JSON_OBJECT.validate_json(received[0].body).get("stream") is not True + assert (streamed.tool_calls, streamed.finish_reason) == ( + (("get_weather", _PARIS), ("get_time", _PARIS)), + "tool_calls", + ) + assert streamed.usage is not None + assert streamed.usage.prompt_tokens_details == _PromptDetails(cached_tokens=64) + assert streamed.usage.completion_tokens_details == _CompletionDetails(reasoning_tokens=1920) + + +def _vllm_frame(identity: str, delta: Mapping[str, JsonValue], finish_reason: str | None = None) -> bytes: + usage: Final = ( + {"usage": {"prompt_tokens": 30, "completion_tokens": 12, "total_tokens": 42}} if finish_reason else {} + ) + body: Final = { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": _VLLM_BACKEND, + "choices": [{"index": 0, "delta": dict(delta), "finish_reason": finish_reason}], + **usage, + } + return b"data: " + json.dumps(body).encode() + b"\n\n" + + +def _vllm_sse(deltas: Sequence[Mapping[str, JsonValue]], finish_reason: str) -> Reply: + identity: Final = f"chatcmpl-audit-{uuid.uuid4().hex}" + return Reply( + content_type="text/event-stream", + chunks=( + _vllm_frame(identity, {"role": "assistant", "content": ""}), + *(_vllm_frame(identity, delta) for delta in deltas), + _vllm_frame(identity, {}, finish_reason), + b"data: [DONE]\n\n", + ), + ) + + +def _is_discovery_probe(request: Request) -> bool: + return (request.method, request.target) == _DISCOVERY_PROBE + + +@contextmanager +def _vllm_server(reply: Reply) -> Generator[Wire, None, None]: + with wire_server( + lambda request: Reply(body=b'{"object":"list","data":[]}') if _is_discovery_probe(request) else reply + ) as wire: + yield wire + + +def _streamed_chats(wire: Wire) -> list[JsonValue]: + posted: Final = [request for request in wire.drain() if not _is_discovery_probe(request)] + assert [(request.method, request.target) for request in posted] == [("POST", "/v1/chat/completions")] + return [_JSON_OBJECT.validate_json(request.body).get("stream") for request in posted] + + +@pytest.mark.parametrize( + "deltas, merged", + [ + pytest.param( + ({"reasoning_content": "R1 "}, {"reasoning_content": "R2 ", "content": "A1 "}, {"content": "A2"}), + "R1 R2 A1 A2", + id="answer-starts-in-the-last-reasoning-chunk", + ), + pytest.param( + ({"reasoning_content": "R ", "content": "A "}, {"content": "B"}), + "R A B", + id="reasoning-and-answer-in-the-first-chunk", + ), + pytest.param( + ({"reasoning_content": "R1 "}, {"reasoning_content": "R2 "}, {"content": "A1 "}, {"content": "A2"}), + "R1 R2 A1 A2", + id="reasoning-and-answer-in-separate-chunks", + ), + ], +) +def test_merged_reasoning_content_closes_the_think_tag_before_the_answer( + gateway: Gateway, deltas: tuple[dict[str, JsonValue], ...], merged: str +) -> None: + with _vllm_server(_vllm_sse(deltas, "stop")) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"hosted_vllm/{_VLLM_BACKEND}", + api_base=wire.url + "/v1", + api_key=_VLLM_KEY, + merge_reasoning_content_in_choices=True, + ) + streamed: Final = _chat_stream( + gateway, {"model": model, "question": f"Hello {uuid.uuid4().hex}", "cache": _NO_CACHE} + ) + + assert _streamed_chats(wire) == [True] + assert (streamed.content, streamed.reasoning, streamed.finish_reason) == (merged, "", "stop") + + +_VLLM_TOOL_CALL: Final[dict[str, JsonValue]] = { + "tool_calls": [ + { + "index": 0, + "id": "call_bridge_weather", + "type": "function", + "function": {"name": "get_weather", "arguments": json.dumps(_PARIS)}, + } + ] +} + + +@pytest.mark.parametrize( + "answer_delta, finish_reason, answer_item", + [ + pytest.param({"content": "It is sunny."}, "stop", "message", id="text"), + pytest.param(_VLLM_TOOL_CALL, "tool_calls", "function_call", id="tool-call"), + ], +) +def test_responses_stream_closes_the_reasoning_item_before_an_answer_from_the_same_chunk( + gateway: Gateway, answer_delta: dict[str, JsonValue], finish_reason: str, answer_item: str +) -> None: + marker: Final = uuid.uuid4().hex + reasoning: Final = f"Bridge reasoning {marker}." + reply: Final = _vllm_sse(({"reasoning_content": reasoning, **answer_delta},), finish_reason) + with _vllm_server(reply) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"hosted_vllm/{_VLLM_BACKEND}", + api_base=wire.url + "/v1", + api_key=_VLLM_KEY, + use_chat_completions_api=True, + ) + events: Final = _raw_responses_events( + gateway, + { + "model": model, + "input": f"Weather in Paris? {marker}", + "tools": [ + { + "type": "function", + "name": "get_weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + } + ], + }, + ) + + lifecycle, summary, text = _responses_reasoning_and_answer(events) + assert _streamed_chats(wire) == [True] + assert lifecycle == [ + ("response.output_item.added", "reasoning", 0), + ("response.output_item.done", "reasoning", 0), + ("response.output_item.added", answer_item, 1), + ("response.output_item.done", answer_item, 1), + ], [event["type"] for event in events] + assert summary == reasoning + assert text == str(answer_delta.get("content", "")) diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index fa649b48a51..93f700b1c2a 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -6143,72 +6143,6 @@ def test_unit_test_perplexity_citations_chunk(): assert "citations" in response_chunk -def test_mock_response_iterator_tool_use(): - """ - Relevant Issue: https://github.com/BerriAI/litellm/issues/7364 - """ - from litellm.llms.bedrock.chat.invoke_handler import MockResponseIterator - from litellm.types.utils import ( - ChatCompletionMessageToolCall, - Choices, - CompletionTokensDetailsWrapper, - Function, - Message, - PromptTokensDetailsWrapper, - Usage, - ) - - litellm.set_verbose = False - response = ModelResponse( - id="chatcmpl-Ai8KRI5vJPZXQ9SQvEJfTVuVqkyEZ", - created=1735081811, - model="o1-2024-12-17", - object="chat.completion", - system_fingerprint="fp_e6d02d4a78", - choices=[ - Choices( - finish_reason="tool_calls", - index=0, - message=Message( - content=None, - role="assistant", - tool_calls=[ - ChatCompletionMessageToolCall( - function=Function( - arguments='{"location":"San Francisco, CA","unit":"fahrenheit"}', - name="get_current_weather", - ), - id="call_BfRX2S7YCKL0BtxbWMl89ZNk", - type="function", - ) - ], - function_call=None, - ), - ) - ], - usage=Usage( - completion_tokens=1955, - prompt_tokens=85, - total_tokens=2040, - completion_tokens_details=CompletionTokensDetailsWrapper( - accepted_prediction_tokens=0, - audio_tokens=0, - reasoning_tokens=1920, - rejected_prediction_tokens=0, - text_tokens=None, - ), - prompt_tokens_details=PromptTokensDetailsWrapper( - audio_tokens=0, cached_tokens=0, text_tokens=None, image_tokens=None - ), - ), - service_tier=None, - ) - completion_stream = MockResponseIterator(model_response=response) - response_chunk = completion_stream._chunk_parser(chunk_data=response) - - assert response_chunk["tool_use"] is not None - - def test_is_delta_empty(): from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.types.utils import Delta diff --git a/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index 2f0a7ad070b..9e0bd5c60b8 100644 --- a/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -287,7 +287,7 @@ def test_translate_streaming_openai_chunk_to_anthropic_thinking_content_block(): assert block_type == "thinking" assert content_block_start == { "type": "thinking", - "thinking": "I need to summar", + "thinking": "", "signature": "", } diff --git a/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_combined_chunk.py b/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_combined_chunk.py index 9a6940ff12b..b73df255cad 100644 --- a/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_combined_chunk.py +++ b/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_combined_chunk.py @@ -12,7 +12,11 @@ content survives. import asyncio import json from types import SimpleNamespace -from typing import AsyncIterator +from typing import AsyncIterator, Final + +import anthropic +import httpx +import pytest from litellm.llms.anthropic.pass_through.adapters.streaming_iterator import ( AnthropicStreamWrapper, @@ -273,3 +277,93 @@ def test_thinking_then_signature_chunk_does_not_crash_stream(): assert "message_stop" in sse assert "Done" in sse + + +def _fake_stream_with_thinking(thinking_blocks: list[dict[str, str]]) -> MockResponseIterator: + model_response: Final = ModelResponse() + model_response.choices = [ + Choices( + index=0, + message=Message( + role="assistant", + content='{"a": 1}', + reasoning_content="".join(block.get("thinking", "") for block in thinking_blocks), + thinking_blocks=thinking_blocks, + ), + finish_reason="stop", + ) + ] + model_response.usage = Usage(prompt_tokens=5, completion_tokens=7, total_tokens=12) + return MockResponseIterator(model_response=model_response) + + +def _anthropic_sdk_final_content( + completion_stream: "MockResponseIterator | AsyncIterator[ModelResponseStream]", +) -> list[dict[str, object]]: + sse: Final = _collect_async(AnthropicStreamWrapper(completion_stream=completion_stream, model="claude")) + transport: Final = httpx.MockTransport( + lambda request: httpx.Response(200, content=sse.encode(), headers={"content-type": "text/event-stream"}) + ) + client: Final = anthropic.Anthropic(api_key="transport-only", http_client=httpx.Client(transport=transport)) + with client.messages.stream( + model="claude", max_tokens=16, messages=[{"role": "user", "content": "hi"}] + ) as stream: + return [block.model_dump(exclude_none=True) for block in stream.get_final_message().content] + + +@pytest.mark.parametrize( + "thinking_blocks", + [ + [ + {"type": "thinking", "thinking": "first", "signature": "SIG_A"}, + {"type": "thinking", "thinking": "second", "signature": "SIG_B"}, + ], + [ + {"type": "redacted_thinking", "data": "OPAQUE"}, + {"type": "thinking", "thinking": "visible", "signature": "SIG_C"}, + ], + [ + {"type": "thinking", "thinking": "unsigned", "signature": ""}, + {"type": "thinking", "thinking": "signed", "signature": "SIG_D"}, + ], + ], + ids=["two_signed_thinking_blocks", "redacted_then_signed_thinking", "unsigned_then_signed_thinking"], +) +def test_fake_stream_gives_each_thinking_block_its_own_content_block(thinking_blocks: list[dict[str, str]]): + content: Final = _anthropic_sdk_final_content(_fake_stream_with_thinking(thinking_blocks)) + + assert content == [*thinking_blocks, {"type": "text", "text": '{"a": 1}'}] + + +def _native_thinking_chunk(thinking: str, index: int, signature: str | None = None) -> ModelResponseStream: + block: Final = {"type": "thinking", "thinking": thinking, "index": index, **({"signature": signature} if signature else {})} + return ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(reasoning_content=thinking or None, thinking_blocks=[block]))] + ) + + +async def _native_stream_with_two_thinking_blocks() -> "AsyncIterator[ModelResponseStream]": + for chunk in ( + _native_thinking_chunk("ab", 0), + _native_thinking_chunk("cd", 0), + _native_thinking_chunk("", 0, "SIG_A"), + _native_thinking_chunk("ef", 1), + _native_thinking_chunk("gh", 1), + _native_thinking_chunk("", 1, "SIG_B"), + ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="hi"))]), + ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(), finish_reason="stop")], + usage=Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2), + ), + ): + yield chunk + + +def test_native_stream_thinking_blocks_keep_their_text_once(): + content: Final = _anthropic_sdk_final_content(_native_stream_with_two_thinking_blocks()) + + assert content == [ + {"type": "thinking", "thinking": "abcd", "signature": "SIG_A"}, + {"type": "thinking", "thinking": "efgh", "signature": "SIG_B"}, + {"type": "text", "text": "hi"}, + ] diff --git a/tests/unit/llms/base_llm/test_base_model_iterator.py b/tests/unit/llms/base_llm/test_base_model_iterator.py index 936bf06bc18..204b458cadc 100644 --- a/tests/unit/llms/base_llm/test_base_model_iterator.py +++ b/tests/unit/llms/base_llm/test_base_model_iterator.py @@ -11,9 +11,9 @@ import httpx import pytest from pydantic import BaseModel -from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator +from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator, convert_model_response_to_streaming from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler -from litellm.types.utils import GenericStreamingChunk +from litellm.types.utils import Choices, GenericStreamingChunk, Message, ModelResponse class TestBaseModelResponseIterator: @@ -314,3 +314,25 @@ def test_close_releases_the_attached_http_response(): iterator.http_response = response iterator.close() assert response.is_closed + + +def test_fake_stream_chunk_indexes_thinking_blocks_by_position_without_touching_the_response(): + """The OpenAI SDK's `.stream()` accumulator merges streamed `thinking_blocks` entries by `index`, so a fake stream + numbers them the way a real stream does, redacted blocks included""" + thinking_blocks: Final = [ + {"type": "thinking", "thinking": "Step 1.", "signature": "sig-a"}, + {"type": "redacted_thinking", "data": "opaque"}, + {"type": "thinking", "thinking": "Step 2.", "signature": "sig-b"}, + ] + response: Final = ModelResponse( + choices=[Choices(message=Message(content="Paris", thinking_blocks=thinking_blocks), finish_reason="stop")] + ) + + chunk: Final = convert_model_response_to_streaming(response) + + assert chunk.choices[0].delta.thinking_blocks == [ + {"type": "thinking", "thinking": "Step 1.", "signature": "sig-a", "index": 0}, + {"type": "redacted_thinking", "data": "opaque", "index": 1}, + {"type": "thinking", "thinking": "Step 2.", "signature": "sig-b", "index": 2}, + ] + assert [block.get("index") for block in response.choices[0].message.thinking_blocks] == [None, None, None] diff --git a/tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py b/tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py index c025a368e51..dec32b3325e 100644 --- a/tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py +++ b/tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py @@ -7,6 +7,7 @@ from __future__ import annotations import asyncio from datetime import datetime, timedelta, timezone +from itertools import chain from typing import Final from unittest.mock import MagicMock, patch @@ -302,3 +303,230 @@ async def test_async_converse_streaming_fails_at_the_request_timeout_not_the_ups def test_sync_converse_streaming_fails_at_the_request_timeout_not_the_upstreams_pace() -> None: with pytest.raises(litellm.Timeout): litellm.completion(client=slow_upstream_sync_client(), **_converse_streaming_kwargs()) + + +RECORDED_REASONING: Final = "Tokyo proper has about 14 million people." +RECORDED_THINKING_TOKENS: Final = 309 +JSON_MODE_THINKING_CONVERSE_RESPONSE: Final = { + "output": { + "message": { + "role": "assistant", + "content": [ + { + "reasoningContent": { + "reasoningText": { + "text": RECORDED_REASONING, + "signature": "EqoBCkgIBxABGAIiQL1", + } + } + }, + { + "toolUse": { + "toolUseId": "tooluse_kZJMlvQmRJ6eAyJE5GIl7Q", + "name": "json_tool_call", + "input": {"city": "Tokyo", "population_millions": 14.0}, + } + }, + ], + } + }, + "stopReason": "tool_use", + "usage": {"inputTokens": 473, "outputTokens": 390, "totalTokens": 863}, + "additionalModelResponseFields": { + "usage": {"output_tokens_details": {"thinking_tokens": RECORDED_THINKING_TOKENS}} + }, +} + + +def _answer_with_json_mode_thinking_response(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json=JSON_MODE_THINKING_CONVERSE_RESPONSE, request=request) + + +def _bedrock_claude_without_native_structured_output() -> str: + return next( + f"bedrock/{model}" + for model, info in litellm.model_cost.items() + if info.get("litellm_provider") == "bedrock_converse" + and model.startswith("us.anthropic.claude") + and info.get("supports_native_structured_output") is False + and info.get("supports_reasoning") is True + ) + + +def _json_mode_thinking_kwargs() -> dict[str, object]: + return { + "model": _bedrock_claude_without_native_structured_output(), + "messages": [{"role": "user", "content": "How many people live in Tokyo, in millions?"}], + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "city_fact", + "schema": { + "type": "object", + "properties": {"city": {"type": "string"}, "population_millions": {"type": "number"}}, + "required": ["city", "population_millions"], + }, + }, + }, + "thinking": {"type": "enabled", "budget_tokens": 1024}, + "aws_access_key_id": "fake", + "aws_secret_access_key": "fake", + "aws_region_name": "us-east-1", + } + + +def _answer_fields(response: ModelResponse) -> tuple[object, ...]: + choice: Final = response.choices[0] + return ( + choice.finish_reason, + choice.message.content, + choice.message.tool_calls, + getattr(choice.message, "reasoning_content", None), + getattr(choice.message, "thinking_blocks", None), + getattr(response.usage.completion_tokens_details, "reasoning_tokens", None), + ) + + +def _streamed_thinking_blocks(chunks: list[litellm.ModelResponseStream]) -> list[object]: + deltas: Final = [chunk.choices[0].delta for chunk in chunks if chunk.choices] + block_lists: Final = [getattr(delta, "thinking_blocks", None) or [] for delta in deltas] + return list(chain.from_iterable(block_lists)) + + +def test_fake_streamed_json_mode_answer_keeps_the_reasoning_of_the_non_streamed_answer() -> None: + """Without native structured output, json_schema plus streaming fakes the stream from one Converse call. That fake + stream dropped the reasoning, the signed thinking blocks and the reasoning tokens, and ended on tool_calls""" + transport: Final = httpx.MockTransport(_answer_with_json_mode_thinking_response) + client: Final = HTTPHandler(client=httpx.Client(transport=transport)) + kwargs: Final = _json_mode_thinking_kwargs() + + answer: Final = litellm.completion(client=client, **kwargs) + chunks: Final = list( + litellm.completion(client=client, stream=True, stream_options={"include_usage": True}, **kwargs) + ) + rebuilt: Final = litellm.stream_chunk_builder(chunks, messages=kwargs["messages"]) + + finish_reason, _, _, reasoning_content, _, reasoning_tokens = _answer_fields(answer) + assert (finish_reason, reasoning_content, reasoning_tokens) == ( + "stop", + RECORDED_REASONING, + RECORDED_THINKING_TOKENS, + ) + assert _answer_fields(rebuilt) == _answer_fields(answer) + assert _streamed_thinking_blocks(chunks) == [ + {**block, "index": index} for index, block in enumerate(answer.choices[0].message.thinking_blocks) + ] + + +@pytest.mark.asyncio +async def test_async_fake_streamed_json_mode_answer_keeps_the_reasoning_of_the_non_streamed_answer() -> None: + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_answer_with_json_mode_thinking_response)) + kwargs: Final = _json_mode_thinking_kwargs() + + answer: Final = await litellm.acompletion(client=client, **kwargs) + stream: Final = await litellm.acompletion( + client=client, stream=True, stream_options={"include_usage": True}, **kwargs + ) + chunks: Final = [chunk async for chunk in stream] + rebuilt: Final = litellm.stream_chunk_builder(chunks, messages=kwargs["messages"]) + + finish_reason, _, _, reasoning_content, _, reasoning_tokens = _answer_fields(answer) + assert (finish_reason, reasoning_content, reasoning_tokens) == ( + "stop", + RECORDED_REASONING, + RECORDED_THINKING_TOKENS, + ) + assert _answer_fields(rebuilt) == _answer_fields(answer) + assert _streamed_thinking_blocks(chunks) == [ + {**block, "index": index} for index, block in enumerate(answer.choices[0].message.thinking_blocks) + ] + + +LLAMA_TEXT_TOOL_CALL_CONVERSE_RESPONSE: Final = { + "output": { + "message": { + "role": "assistant", + "content": [{"text": '{"type": "function", "name": "get_weather", "parameters": {"city": "Paris"}}'}], + } + }, + "stopReason": "end_turn", + "usage": {"inputTokens": 175, "outputTokens": 23, "totalTokens": 198}, +} + + +def _answer_with_llama_text_tool_call(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json=LLAMA_TEXT_TOOL_CALL_CONVERSE_RESPONSE, request=request) + + +def _llama_tool_kwargs() -> dict[str, object]: + return { + "model": "bedrock/us.meta.llama3-3-70b-instruct-v1:0", + "messages": [{"role": "user", "content": "What is the weather in Paris right now?"}], + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + } + ], + "aws_access_key_id": "fake", + "aws_secret_access_key": "fake", + "aws_region_name": "us-east-1", + } + + +def _tool_answer(response: ModelResponse) -> tuple[object, ...]: + choice: Final = response.choices[0] + return ( + choice.finish_reason, + choice.message.content, + [(call.function.name, call.function.arguments) for call in choice.message.tool_calls or []], + ) + + +def test_fake_streamed_text_tool_call_arrives_as_the_tool_call_of_the_non_streamed_answer() -> None: + client: Final = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_answer_with_llama_text_tool_call))) + kwargs: Final = _llama_tool_kwargs() + + answer: Final = litellm.completion(client=client, **kwargs) + chunks: Final = list(litellm.completion(client=client, stream=True, **kwargs)) + + assert _tool_answer(answer) == ("tool_calls", None, [("get_weather", '{"city": "Paris"}')]) + assert _tool_answer(litellm.stream_chunk_builder(chunks, messages=kwargs["messages"])) == _tool_answer(answer) + + +@pytest.mark.asyncio +async def test_async_fake_streamed_text_tool_call_arrives_as_the_tool_call_of_the_non_streamed_answer() -> None: + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_answer_with_llama_text_tool_call)) + kwargs: Final = _llama_tool_kwargs() + + answer: Final = await litellm.acompletion(client=client, **kwargs) + stream: Final = await litellm.acompletion(client=client, stream=True, **kwargs) + chunks: Final = [chunk async for chunk in stream] + + assert _tool_answer(answer) == ("tool_calls", None, [("get_weather", '{"city": "Paris"}')]) + assert _tool_answer(litellm.stream_chunk_builder(chunks, messages=kwargs["messages"])) == _tool_answer(answer) + + +def test_fake_streamed_json_mode_answer_closes_merged_reasoning_before_the_answer() -> None: + """A fake stream sends reasoning and answer in one chunk, so merging reasoning into content must open and close + the think block around the reasoning in that chunk, before the answer""" + client: Final = HTTPHandler( + client=httpx.Client(transport=httpx.MockTransport(_answer_with_json_mode_thinking_response)) + ) + kwargs: Final = _json_mode_thinking_kwargs() + + answer: Final = litellm.completion(client=client, **kwargs) + chunks: Final = list( + litellm.completion(client=client, stream=True, merge_reasoning_content_in_choices=True, **kwargs) + ) + + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == ( + f"{answer.choices[0].message.reasoning_content}{answer.choices[0].message.content}" + ) diff --git a/tests/unit/llms/openai/test_openai.py b/tests/unit/llms/openai/test_openai.py index e37fd4b6ae5..d7ca9319f3a 100644 --- a/tests/unit/llms/openai/test_openai.py +++ b/tests/unit/llms/openai/test_openai.py @@ -10,7 +10,7 @@ from unittest.mock import Mock, patch import httpx import pytest import respx -from openai import AsyncOpenAI, OpenAI +from openai import AsyncOpenAI, AzureOpenAI, OpenAI from openai._types import Response as SDKResponse from pydantic import BaseModel, ConfigDict, JsonValue @@ -221,6 +221,79 @@ async def test_acompletion_returns_prompt_cache_details_over_injected_transport( assert prompt_tokens_details.cached_tokens == 7 +_REASONING_TOOL_CALL_REPLY: Final = { + "id": "chatcmpl-fake-stream", + "object": "chat.completion", + "created": 0, + "model": "o4-mini", + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_weather", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + }, + { + "id": "call_time", + "type": "function", + "function": {"name": "get_time", "arguments": '{"city": "Paris"}'}, + }, + ], + }, + } + ], + "usage": { + "prompt_tokens": 85, + "completion_tokens": 1955, + "total_tokens": 2040, + "prompt_tokens_details": {"cached_tokens": 64}, + "completion_tokens_details": {"reasoning_tokens": 1920}, + }, +} + + +def test_azure_o_series_fake_stream_keeps_every_tool_call_and_token_detail(): + request_bodies: Final[list[dict[str, object]]] = [] + + def respond(request: httpx.Request) -> httpx.Response: + request_bodies.append(json.loads(request.content)) + return httpx.Response(200, json=_REASONING_TOOL_CALL_REPLY) + + client: Final = AzureOpenAI( + api_key="transport-only", + api_version="2025-03-01-preview", + azure_endpoint="https://fake-stream.openai.azure.com", + http_client=httpx.Client(transport=httpx.MockTransport(respond)), + ) + chunks: Final = list( + litellm.completion( + model="azure/o_series/reasoning-deployment", + api_key="transport-only", + api_version="2025-03-01-preview", + api_base="https://fake-stream.openai.azure.com", + client=client, + messages=[{"role": "user", "content": "weather and time in Paris?"}], + stream=True, + stream_options={"include_usage": True}, + num_retries=0, + ) + ) + rebuilt: Final = litellm.stream_chunk_builder(chunks) + + assert [body.get("stream") for body in request_bodies] == [None] + assert rebuilt is not None + assert rebuilt.choices[0].finish_reason == "tool_calls" + assert [call.function.name for call in rebuilt.choices[0].message.tool_calls] == ["get_weather", "get_time"] + assert rebuilt.usage.prompt_tokens_details.cached_tokens == 64 + assert rebuilt.usage.completion_tokens_details.reasoning_tokens == 1920 + + @respx.mock def test_completion_defaults_missing_message_role_and_preserves_omitted_fields() -> None: route: Final = respx.post("https://api.openai.com/v1/chat/completions").mock( diff --git a/tests/unit/llms/openai_like/chat/test_openai_like_chat_handler.py b/tests/unit/llms/openai_like/chat/test_openai_like_chat_handler.py new file mode 100644 index 00000000000..75b60b2dd6c --- /dev/null +++ b/tests/unit/llms/openai_like/chat/test_openai_like_chat_handler.py @@ -0,0 +1,75 @@ +import json +from typing import Final + +from unittest.mock import Mock + +import httpx + +from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.openai_like.chat.handler import make_sync_call + +_FAKE_STREAMED_REPLY: Final = { + "id": "chatcmpl-fake-stream", + "object": "chat.completion", + "created": 0, + "model": "openai/gpt-oss-120b", + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": None, + "reasoning_content": "Look up both values for Paris.", + "tool_calls": [ + { + "id": "call_weather", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + }, + { + "id": "call_time", + "type": "function", + "function": {"name": "get_time", "arguments": '{"city": "Paris"}'}, + }, + ], + }, + } + ], + "usage": { + "prompt_tokens": 85, + "completion_tokens": 1955, + "total_tokens": 2040, + "prompt_tokens_details": {"cached_tokens": 64}, + "completion_tokens_details": {"reasoning_tokens": 1920}, + }, +} + + +def test_fake_stream_keeps_reasoning_every_tool_call_and_token_detail(): + request_bodies: Final[list[dict[str, object]]] = [] + + def respond(request: httpx.Request) -> httpx.Response: + request_bodies.append(json.loads(request.content)) + return httpx.Response(200, json=_FAKE_STREAMED_REPLY) + + chunks: Final = list( + make_sync_call( + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))), + api_base="https://openai-like.example/v1/chat/completions", + headers={}, + data=json.dumps({"model": "openai/gpt-oss-120b", "messages": []}), + model="openai/gpt-oss-120b", + messages=[], + logging_obj=Mock(), + fake_stream=True, + ) + ) + + assert len(request_bodies) == 1 + assert len(chunks) == 1 + delta: Final = chunks[0].choices[0].delta + assert delta.reasoning_content == "Look up both values for Paris." + assert [call.function.name for call in delta.tool_calls] == ["get_weather", "get_time"] + assert chunks[0].usage.prompt_tokens_details.cached_tokens == 64 + assert chunks[0].usage.completion_tokens_details.reasoning_tokens == 1920 diff --git a/tests/unit/llms/snowflake/chat/test_snowflake_chat_transformation.py b/tests/unit/llms/snowflake/chat/test_snowflake_chat_transformation.py index 5687a319f06..f9d12b7eeea 100644 --- a/tests/unit/llms/snowflake/chat/test_snowflake_chat_transformation.py +++ b/tests/unit/llms/snowflake/chat/test_snowflake_chat_transformation.py @@ -8,7 +8,7 @@ import os import copy import uuid import json -from typing import Any, Dict, List +from typing import Any, Dict, Final, List from unittest.mock import AsyncMock, patch, Mock, MagicMock @@ -595,6 +595,26 @@ class TestSnowflakeCortexClaudeFixes: assert blocks[0] == {"type": "thinking", "thinking": "391", "signature": "Eto"} assert [b["type"] for b in blocks] == ["thinking", "tool_use"] + def test_signed_thinking_blocks_replay_without_their_stream_index(self): + body: Final = self._transform( + [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": "391", + "thinking_blocks": [ + {"type": "thinking", "thinking": "391", "signature": "Eto", "index": 0}, + {"type": "redacted_thinking", "data": "opaque", "index": 1}, + ], + }, + ] + ) + blocks: Final = body["messages"][1]["content"] + assert blocks[:2] == [ + {"type": "thinking", "thinking": "391", "signature": "Eto"}, + {"type": "redacted_thinking", "data": "opaque"}, + ] + def test_signed_thinking_blocks_lead_a_plain_text_assistant_turn(self): """A thinking response without a tool call must also round-trip on the next request.""" body = self._transform( diff --git a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 4117c303940..1a3430224d1 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -11,6 +11,7 @@ spend tracking stores, so a follow-up previous_response_id still finds the conve """ import json +from collections.abc import Callable from typing import Final from unittest.mock import AsyncMock, MagicMock @@ -978,6 +979,44 @@ def _reasoning_chunk(reasoning: str, finish_reason: str | None = None) -> ModelR ) +def _with_reasoning(chunk: ModelResponseStream, reasoning: str) -> ModelResponseStream: + choice: Final = chunk.choices[0] + delta: Final = Delta(**{**choice.delta.model_dump(), "reasoning_content": reasoning}) + return chunk.model_copy( + update={"choices": [StreamingChoices(index=0, delta=delta, finish_reason=choice.finish_reason)]} + ) + + +def _reasoning_then_text_chunks() -> tuple[ModelResponseStream, ...]: + return (_reasoning_chunk("let me think"), _chunk("Hello"), _chunk("!", finish_reason="stop")) + + +def _reasoning_and_text_in_one_chunk() -> tuple[ModelResponseStream, ...]: + return (_with_reasoning(_chunk("Hello!", finish_reason="stop"), "let me think"),) + + +def _signed_thinking_and_text_in_one_chunk() -> tuple[ModelResponseStream, ...]: + chunk: Final = _chunk("Hello!", finish_reason="stop") + signed_block: Final = {"type": "thinking", "thinking": "", "signature": "sig_only"} + delta: Final = Delta(**{**chunk.choices[0].delta.model_dump(), "thinking_blocks": [signed_block]}) + return ( + chunk.model_copy(update={"choices": [StreamingChoices(index=0, delta=delta, finish_reason="stop")]}), + ) + + +StreamChunks = Callable[[], tuple[ModelResponseStream, ...]] + +REASONING_THEN_TEXT_STREAMS: Final = pytest.mark.parametrize( + "stream_chunks", + [_reasoning_then_text_chunks, _reasoning_and_text_in_one_chunk], + ids=["reasoning_then_text_chunks", "reasoning_and_text_in_one_chunk"], +) + + +def _event_text(events: list[BaseLiteLLMOpenAIResponseObject], event_type: ResponsesAPIStreamEvents) -> str: + return "".join(event.delta for event in events if getattr(event, "type", None) == event_type) + + def _annotation_only_chunk() -> ModelResponseStream: citation: Final = {"start_index": 0, "end_index": 2, "url": "https://example.com", "title": "Example"} return ModelResponseStream( @@ -1134,16 +1173,13 @@ async def test_signature_only_thinking_streams_a_replayable_reasoning_item(sync_ assert json.loads(reasoning_items[0].encrypted_content)[0]["signature"] == "sig_only" +@REASONING_THEN_TEXT_STREAMS @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio -async def test_reasoning_then_text_announces_message_item_before_text_events(sync_mode: bool): - iterator: Final = _build_iterator( - [ - _reasoning_chunk("let me think"), - _chunk("Hello"), - _chunk("!", finish_reason="stop"), - ] - ) +async def test_reasoning_then_text_announces_message_item_before_text_events( + stream_chunks: StreamChunks, sync_mode: bool +): + iterator: Final = _build_iterator(stream_chunks()) events: Final = await _collect_events(iterator, sync_mode) @@ -1171,23 +1207,22 @@ async def test_reasoning_then_text_announces_message_item_before_text_events(syn elif event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE and _is_message_item(event): assert event.item.id in announced_message_ids assert saw_text_delta - assert "".join( - event.delta for event in events if getattr(event, "type", None) == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA - ) == "Hello!" + assert _event_text(events, ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA) == "Hello!" + assert _event_text(events, ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DELTA) == "let me think" assert announced_indexes_by_item_type["message"] != announced_indexes_by_item_type["reasoning"] +@pytest.mark.parametrize( + "stream_chunks", + [_reasoning_then_text_chunks, _reasoning_and_text_in_one_chunk, _signed_thinking_and_text_in_one_chunk], + ids=["reasoning_then_text_chunks", "reasoning_and_text_in_one_chunk", "signed_thinking_and_text_in_one_chunk"], +) +@pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio -async def test_reasoning_item_closes_before_message_item_opens(): - iterator: Final = _build_iterator( - [ - _reasoning_chunk("let me think"), - _chunk("Hello"), - _chunk("!", finish_reason="stop"), - ] - ) +async def test_reasoning_item_closes_before_message_item_opens(stream_chunks: StreamChunks, sync_mode: bool): + iterator: Final = _build_iterator(stream_chunks()) - events: Final = await _collect_events(iterator, sync_mode=False) + events: Final = await _collect_events(iterator, sync_mode) item_lifecycle: Final = [ (event.type, event.item.type) @@ -1201,6 +1236,7 @@ async def test_reasoning_item_closes_before_message_item_opens(): (ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, "message"), (ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, "message"), ] + assert _event_text(events, ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA) == "Hello!" @pytest.mark.parametrize("sync_mode", [True, False]) @@ -1231,6 +1267,27 @@ async def test_tool_then_reasoning_then_text_gives_message_its_own_output_index( assert len(output_indexes_by_item_id) == len(set(output_indexes_by_item_id.values())) +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_reasoning_item_closes_before_a_tool_call_from_the_same_chunk_opens(sync_mode: bool): + iterator: Final = _build_iterator([_with_reasoning(_tool_call_chunk(finish_reason="tool_calls"), "let me think")]) + + events: Final = await _collect_events(iterator, sync_mode) + + item_lifecycle: Final = [ + (event.type, event.item.type) + for event in events + if getattr(event, "type", None) + in (ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE) + ] + assert item_lifecycle == [ + (ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, "reasoning"), + (ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, "reasoning"), + (ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, "function_call"), + (ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, "function_call"), + ] + + @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio async def test_tool_only_stream_puts_the_function_call_at_output_index_zero(sync_mode: bool): From ee4860dd84305dd07edd492247db853151ef90c3 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 16 Oct 2026 10:31:41 -0700 Subject: [PATCH 06/58] test(cost): add tier 1 cost tracking cases for gpt-5.6-terra, gemini-4-argon and gemini-3.1-pro-preview (#46931) * test(cost): add tier 1 cost tracking cases for gpt-5.6-terra, gemini-4-argon and gemini-3.1-pro-preview Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(cost): add audit 3 shape cases for vertex_ai and azure_ai claude Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(cost): mark azure gpt-5.6-terra priority case as known bug Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../cost_calculation/cases/anthropic.py | 16 + .../cost_calculation/cases/azure.py | 175 +++++++++- .../cost_calculation/cases/gemini.py | 89 ++++- .../cost_calculation/cases/misc.py | 247 ++++++++++++- .../cost_calculation/cases/openai.py | 188 ++++++++++ .../cost_calculation/cases/vertex_ai.py | 327 +++++++++++++++++- .../integration/cost_calculation/cost_map.py | 145 ++++++++ 7 files changed, 1180 insertions(+), 7 deletions(-) diff --git a/tests/integration/cost_calculation/cases/anthropic.py b/tests/integration/cost_calculation/cases/anthropic.py index 3d74ac4e2f3..a0c81d8f06d 100644 --- a/tests/integration/cost_calculation/cases/anthropic.py +++ b/tests/integration/cost_calculation/cases/anthropic.py @@ -4111,6 +4111,7 @@ CLAUDE_OPUS_5_5_INPUT_TEXT: Final = CostTrackingTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, }, ), @@ -4170,6 +4171,7 @@ CLAUDE_SONNET_5_5_INPUT_TEXT: Final = CostTrackingTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -4232,6 +4234,7 @@ CLAUDE_SONNET_5_5_STREAM: Final = CostTrackingTestCase( "output_tokens": 1, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -4314,6 +4317,7 @@ CLAUDE_SONNET_5_5_CACHE_READ: Final = CostTrackingTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -4566,6 +4570,7 @@ CLAUDE_HAIKU_5_5_INPUT_TEXT: Final = CostTrackingTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -4628,6 +4633,7 @@ CLAUDE_HAIKU_5_5_STREAM: Final = CostTrackingTestCase( "output_tokens": 1, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -4710,6 +4716,7 @@ CLAUDE_HAIKU_5_5_CACHE_READ: Final = CostTrackingTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -5325,6 +5332,7 @@ CLAUDE_FABLE_5_1_INPUT_TEXT: Final = CostTrackingTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -5387,6 +5395,7 @@ CLAUDE_FABLE_5_1_STREAM: Final = CostTrackingTestCase( "output_tokens": 1, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -5469,6 +5478,7 @@ CLAUDE_FABLE_5_1_CACHE_READ: Final = CostTrackingTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -5720,6 +5730,7 @@ CLAUDE_OPUS_5_5_STREAM: Final = CostTrackingTestCase( "output_tokens": 1, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -5802,6 +5813,7 @@ CLAUDE_OPUS_5_5_CACHE_READ: Final = CostTrackingTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -6141,6 +6153,7 @@ CLAUDE_SONNET_5_5_THINKING: Final = CostTrackingTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -6294,6 +6307,7 @@ CLAUDE_HAIKU_5_5_THINKING: Final = CostTrackingTestCase( "output_tokens_details": {"thinking_tokens": 3480}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -6593,6 +6607,7 @@ CLAUDE_FABLE_5_1_THINKING: Final = CostTrackingTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -6746,6 +6761,7 @@ CLAUDE_OPUS_5_5_THINKING: Final = CostTrackingTestCase( "output_tokens_details": {"thinking_tokens": 3480}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, diff --git a/tests/integration/cost_calculation/cases/azure.py b/tests/integration/cost_calculation/cases/azure.py index 14c9647dc08..3df038d1e0e 100644 --- a/tests/integration/cost_calculation/cases/azure.py +++ b/tests/integration/cost_calculation/cases/azure.py @@ -2,7 +2,7 @@ PARITY pairs the streamed case with its plain twin where both bill the same row.""" -from typing import Final +from typing import Final, cast from integration.cost_calculation.cost_tracking_case import ( with_response_usage, @@ -23,6 +23,20 @@ from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParit from integration.cost_calculation.stream_parity.case import StreamParityTestCase, sse_frames +def _response_for_model(response: JsonResponse, model: str) -> JsonResponse: + return response.model_copy(update={"body": {**response.body, "model": model}}) + + +def _stream_for_model(response: SseResponse, model: str) -> SseResponse: + return response.model_copy( + update={ + "frames": tuple( + frame.replace("gpt-6-luna-2026-09-22", model).replace("gpt-6-luna", model) for frame in response.frames + ) + } + ) + + # azure/gpt-5.4-mini AZURE_GPT_5_4_MINI_INPUT_TEXT: Final = CostTrackingTestCase( name="azure-gpt-5.4-mini-input_text", @@ -6174,6 +6188,152 @@ AZURE_GPT_6_LUNA_CACHE_WRITE_5M: Final = with_response_usage( "content_filter_result": {"error": {"code": "content_filter_error", "message": "The contents are not filtered"}} }, ) + +AZURE_GPT_5_6_TERRA_INPUT_TEXT: Final = AZURE_GPT_6_LUNA_INPUT_TEXT.model_copy( + update={ + "name": "azure-gpt-5.6-terra-input_text", + "model": "azure/gpt-5.6-terra", + "response": _response_for_model( + cast(JsonResponse, AZURE_GPT_6_LUNA_INPUT_TEXT.response), "gpt-5.6-terra-2026-07-09" + ), + "expected": ExactExpected( + spend=0.0094864, + input_cost=0.004048, + output_cost=0.0054384, + prompt_tokens=1840, + completion_tokens=412, + ), + } +) +AZURE_GPT_5_6_TERRA_STREAM: Final = AZURE_GPT_6_LUNA_STREAM.model_copy( + update={ + "name": "azure-gpt-5.6-terra-stream", + "model": "azure/gpt-5.6-terra", + "response": _stream_for_model( + cast(SseResponse, AZURE_GPT_6_LUNA_STREAM.response), "gpt-5.6-terra-2026-07-09" + ), + "expected": ExactExpected( + spend=0.0094864, + input_cost=0.004048, + output_cost=0.0054384, + prompt_tokens=1840, + completion_tokens=412, + ), + } +) +AZURE_GPT_5_6_TERRA_SERVICE_TIER_FLEX: Final = AZURE_GPT_6_LUNA_SERVICE_TIER_FLEX.model_copy( + update={ + "name": "azure-gpt-5.6-terra-service_tier_flex", + "model": "azure/gpt-5.6-terra", + "response": _response_for_model( + cast(JsonResponse, AZURE_GPT_6_LUNA_SERVICE_TIER_FLEX.response), "gpt-5.6-terra-2026-07-09" + ), + "expected": ExactExpected( + spend=0.0047432, + input_cost=0.002024, + output_cost=0.0027192, + prompt_tokens=1840, + completion_tokens=412, + ), + } +) +AZURE_GPT_5_6_TERRA_SERVICE_TIER_PRIORITY: Final = AZURE_GPT_6_LUNA_SERVICE_TIER_PRIORITY.model_copy( + update={ + "name": "azure-gpt-5.6-terra-service_tier_priority", + "model": "azure/gpt-5.6-terra", + "response": _response_for_model( + cast(JsonResponse, AZURE_GPT_6_LUNA_SERVICE_TIER_PRIORITY.response), "gpt-5.6-terra-2026-07-09" + ), + "expected": ExactExpected( + spend=0.0189728, + input_cost=0.008096, + output_cost=0.0108768, + prompt_tokens=1840, + completion_tokens=412, + ), + "known_bug": "Azure GPT Terra priority service tier is billed at the base rate", + } +) +AZURE_GPT_5_6_TERRA_MESSAGES_BRIDGED_JSON: Final = AZURE_GPT_6_LUNA_MESSAGES_BRIDGED_JSON.model_copy( + update={ + "name": "azure-gpt-5.6-terra-messages_bridged_json", + "model": "azure/gpt-5.6-terra", + "response": _response_for_model( + cast(JsonResponse, AZURE_GPT_6_LUNA_MESSAGES_BRIDGED_JSON.response), "gpt-5.6-terra-2026-07-09" + ), + "expected": ExactExpected( + spend=0.0094864, + input_cost=0.004048, + output_cost=0.0054384, + prompt_tokens=1840, + completion_tokens=412, + ), + } +) +AZURE_GPT_5_6_TERRA_CACHE_READ: Final = AZURE_GPT_6_LUNA_CACHE_READ.model_copy( + update={ + "name": "azure/gpt-5.6-terra-cache_read", + "model": "azure/gpt-5.6-terra", + "response": _response_for_model( + cast(JsonResponse, AZURE_GPT_6_LUNA_CACHE_READ.response), "gpt-5.6-terra-2026-07-09" + ), + "expected": ExactExpected( + spend=0.00912736, + input_cost=0.00411136, + output_cost=0.005016, + prompt_tokens=12928, + completion_tokens=380, + ), + } +) +AZURE_GPT_5_6_TERRA_CACHE_WRITE_5M: Final = AZURE_GPT_6_LUNA_CACHE_WRITE_5M.model_copy( + update={ + "name": "azure/gpt-5.6-terra-cache_write_5m", + "model": "azure/gpt-5.6-terra", + "response": _response_for_model( + cast(JsonResponse, AZURE_GPT_6_LUNA_CACHE_WRITE_5M.response), "gpt-5.6-terra-2026-07-09" + ), + "expected": ExactExpected( + spend=0.040216, + input_cost=0.0352, + output_cost=0.005016, + prompt_tokens=12928, + completion_tokens=380, + ), + } +) +AZURE_GPT_5_6_TERRA_REASONING: Final = AZURE_GPT_6_LUNA_REASONING.model_copy( + update={ + "name": "azure/gpt-5.6-terra-reasoning", + "model": "azure/gpt-5.6-terra", + "response": _response_for_model( + cast(JsonResponse, AZURE_GPT_6_LUNA_REASONING.response), "gpt-5.6-terra-2026-07-09" + ), + "expected": ExactExpected( + spend=0.056056, + input_cost=0.002728, + output_cost=0.053328, + prompt_tokens=1240, + completion_tokens=4040, + ), + } +) +AZURE_GPT_5_6_TERRA_TOOL_CALL: Final = AZURE_GPT_6_LUNA_TOOL_CALL.model_copy( + update={ + "name": "azure/gpt-5.6-terra-tool_call", + "model": "azure/gpt-5.6-terra", + "response": _response_for_model( + cast(JsonResponse, AZURE_GPT_6_LUNA_TOOL_CALL.response), "gpt-5.6-terra-2026-07-09" + ), + "expected": ExactExpected( + spend=0.0094864, + input_cost=0.004048, + output_cost=0.0054384, + prompt_tokens=1840, + completion_tokens=412, + ), + } +) AZURE_GPT_6_ASTRA_CACHE_WRITE_5M: Final = with_response_usage( AZURE_GPT_6_ASTRA_CACHE_READ, "azure/gpt-6-astra-cache_write_5m", @@ -6304,6 +6464,15 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( AZURE_GPT_6_1_SOL_CACHE_WRITE_5M, AZURE_GPT_6_LUNA_CACHE_WRITE_5M, AZURE_GPT_6_ASTRA_CACHE_WRITE_5M, + AZURE_GPT_5_6_TERRA_INPUT_TEXT, + AZURE_GPT_5_6_TERRA_STREAM, + AZURE_GPT_5_6_TERRA_SERVICE_TIER_FLEX, + AZURE_GPT_5_6_TERRA_SERVICE_TIER_PRIORITY, + AZURE_GPT_5_6_TERRA_MESSAGES_BRIDGED_JSON, + AZURE_GPT_5_6_TERRA_CACHE_READ, + AZURE_GPT_5_6_TERRA_CACHE_WRITE_5M, + AZURE_GPT_5_6_TERRA_REASONING, + AZURE_GPT_5_6_TERRA_TOOL_CALL, ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( @@ -6316,6 +6485,7 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( StreamParityTestCase(plain=AZURE_GPT_6_LUNA_INPUT_TEXT, streamed=AZURE_GPT_6_LUNA_STREAM), StreamParityTestCase(plain=AZURE_GPT_6_ASTRA_INPUT_TEXT, streamed=AZURE_GPT_6_ASTRA_STREAM), StreamParityTestCase(plain=AZURE_GPT_6_1_SOL_INPUT_TEXT, streamed=AZURE_GPT_6_1_SOL_STREAM), + StreamParityTestCase(plain=AZURE_GPT_5_6_TERRA_INPUT_TEXT, streamed=AZURE_GPT_5_6_TERRA_STREAM), ) BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( @@ -6327,4 +6497,7 @@ BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( BridgeCostParityTestCase(native=AZURE_GPT_6_LUNA_INPUT_TEXT, bridged=AZURE_GPT_6_LUNA_MESSAGES_BRIDGED_JSON), BridgeCostParityTestCase(native=AZURE_GPT_6_ASTRA_INPUT_TEXT, bridged=AZURE_GPT_6_ASTRA_MESSAGES_BRIDGED_JSON), BridgeCostParityTestCase(native=AZURE_GPT_6_1_SOL_INPUT_TEXT, bridged=AZURE_GPT_6_1_SOL_MESSAGES_BRIDGED_JSON), + BridgeCostParityTestCase( + native=AZURE_GPT_5_6_TERRA_INPUT_TEXT, bridged=AZURE_GPT_5_6_TERRA_MESSAGES_BRIDGED_JSON + ), ) diff --git a/tests/integration/cost_calculation/cases/gemini.py b/tests/integration/cost_calculation/cases/gemini.py index a4538b87db3..220ce2023be 100644 --- a/tests/integration/cost_calculation/cases/gemini.py +++ b/tests/integration/cost_calculation/cases/gemini.py @@ -2,7 +2,7 @@ PARITY pairs the streamed case with its plain twin where both bill the same row.""" -from typing import Final +from typing import Final, cast from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParityTestCase, bridge_pairs from integration.cost_calculation.cost_tracking_case import ( @@ -16,6 +16,32 @@ from integration.cost_calculation.cost_tracking_case import ( ) from integration.cost_calculation.stream_parity.case import StreamParityTestCase, sse_frames + +def _gemini_case_response(response: JsonResponse | SseResponse, model_version: str) -> JsonResponse | SseResponse: + match response: + case JsonResponse(): + return response.model_copy(update={"body": {**response.body, "modelVersion": model_version}}) + case SseResponse(): + return response.model_copy( + update={ + "frames": tuple(frame.replace("gemini-3.8-flash", model_version) for frame in response.frames) + } + ) + + +def _gemini_case( + sibling: CostTrackingTestCase, name: str, model: str, model_version: str, expected: ExactExpected +) -> CostTrackingTestCase: + return sibling.model_copy( + update={ + "name": name, + "model": model, + "response": _gemini_case_response(cast(JsonResponse | SseResponse, sibling.response), model_version), + "expected": expected, + } + ) + + # gemini/gemini-3.1-pro GEMINI_GEMINI_3_1_PRO_INPUT_TEXT: Final = CostTrackingTestCase( name="gemini-gemini-3.1-pro-input_text", @@ -3543,6 +3569,51 @@ GEMINI_GEMINI_3_1_PRO_CACHE_CREATION: Final = CostTrackingTestCase( ), ) +GEMINI_GEMINI_4_ARGON_INPUT_TEXT: Final = _gemini_case( + GEMINI_GEMINI_3_8_FLASH_INPUT_TEXT, "gemini-gemini-4-argon-input_text", "gemini/gemini-4-argon", + "gemini-4-argon", ExactExpected(spend=0.00858, input_cost=0.004048, output_cost=0.004532, prompt_tokens=1840, completion_tokens=412) +) +GEMINI_GEMINI_4_ARGON_STREAM: Final = _gemini_case( + GEMINI_GEMINI_3_8_FLASH_STREAM, "gemini-gemini-4-argon-stream", "gemini/gemini-4-argon", + "gemini-4-argon", ExactExpected(spend=0.00858, input_cost=0.004048, output_cost=0.004532, prompt_tokens=1840, completion_tokens=412) +) +GEMINI_GEMINI_4_ARGON_CACHE_READ: Final = _gemini_case( + GEMINI_GEMINI_3_8_FLASH_CACHE_READ, "gemini-gemini-4-argon-cache_read", "gemini/gemini-4-argon", + "gemini-4-argon", ExactExpected(spend=0.00693968, input_cost=0.00275968, output_cost=0.00418, prompt_tokens=12928, completion_tokens=380) +) +GEMINI_GEMINI_4_ARGON_REASONING: Final = _gemini_case( + GEMINI_GEMINI_3_8_FLASH_REASONING, "gemini-gemini-4-argon-reasoning", "gemini/gemini-4-argon", + "gemini-4-argon", ExactExpected(spend=0.047168, input_cost=0.002728, output_cost=0.04444, prompt_tokens=1240, completion_tokens=4040) +) +GEMINI_GEMINI_4_ARGON_TOOL_CALL: Final = _gemini_case( + GEMINI_GEMINI_3_8_FLASH_TOOL_CALL, "gemini-gemini-4-argon-tool_call", "gemini/gemini-4-argon", + "gemini-4-argon", ExactExpected(spend=0.00858, input_cost=0.004048, output_cost=0.004532, prompt_tokens=1840, completion_tokens=412) +) +GEMINI_GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT: Final = _gemini_case( + GEMINI_GEMINI_3_8_FLASH_INPUT_TEXT, "gemini-gemini-3.1-pro-preview-input_text", + "gemini/gemini-3.1-pro-preview", "gemini-3.1-pro-preview", + ExactExpected(spend=0.0112112, input_cost=0.004784, output_cost=0.0064272, prompt_tokens=1840, completion_tokens=412) +) +GEMINI_GEMINI_3_1_PRO_PREVIEW_STREAM: Final = _gemini_case( + GEMINI_GEMINI_3_8_FLASH_STREAM, "gemini-gemini-3.1-pro-preview-stream", + "gemini/gemini-3.1-pro-preview", "gemini-3.1-pro-preview", + ExactExpected(spend=0.0112112, input_cost=0.004784, output_cost=0.0064272, prompt_tokens=1840, completion_tokens=412) +) +GEMINI_GEMINI_3_1_PRO_PREVIEW_CACHE_READ: Final = _gemini_case( + GEMINI_GEMINI_3_8_FLASH_CACHE_READ, "gemini-gemini-3.1-pro-preview-cache_read", + "gemini/gemini-3.1-pro-preview", "gemini-3.1-pro-preview", + ExactExpected(spend=0.01078688, input_cost=0.00485888, output_cost=0.005928, prompt_tokens=12928, completion_tokens=380) +) +GEMINI_GEMINI_3_1_PRO_PREVIEW_REASONING: Final = _gemini_case( + GEMINI_GEMINI_3_8_FLASH_REASONING, "gemini-gemini-3.1-pro-preview-reasoning", + "gemini/gemini-3.1-pro-preview", "gemini-3.1-pro-preview", + ExactExpected(spend=0.066248, input_cost=0.003224, output_cost=0.063024, prompt_tokens=1240, completion_tokens=4040) +) +GEMINI_GEMINI_3_1_PRO_PREVIEW_TOOL_CALL: Final = _gemini_case( + GEMINI_GEMINI_3_8_FLASH_TOOL_CALL, "gemini-gemini-3.1-pro-preview-tool_call", + "gemini/gemini-3.1-pro-preview", "gemini-3.1-pro-preview", + ExactExpected(spend=0.0112112, input_cost=0.004784, output_cost=0.0064272, prompt_tokens=1840, completion_tokens=412) +) CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GEMINI_GEMINI_3_1_PRO_INPUT_TEXT, @@ -3605,6 +3676,16 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GEMINI_NANO_BANANA_2_1_WEB_SEARCH_SINGLE_QUERY, GEMINI_NANO_BANANA_2_1_WEB_SEARCH_MULTIPLE_QUERIES, GEMINI_GEMINI_3_1_PRO_CACHE_CREATION, + GEMINI_GEMINI_4_ARGON_INPUT_TEXT, + GEMINI_GEMINI_4_ARGON_STREAM, + GEMINI_GEMINI_4_ARGON_CACHE_READ, + GEMINI_GEMINI_4_ARGON_REASONING, + GEMINI_GEMINI_4_ARGON_TOOL_CALL, + GEMINI_GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT, + GEMINI_GEMINI_3_1_PRO_PREVIEW_STREAM, + GEMINI_GEMINI_3_1_PRO_PREVIEW_CACHE_READ, + GEMINI_GEMINI_3_1_PRO_PREVIEW_REASONING, + GEMINI_GEMINI_3_1_PRO_PREVIEW_TOOL_CALL, ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( @@ -3626,6 +3707,10 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( streamed=GEMINI_GEMINI_3_8_FLASH_STREAM_RESPONSE_MODEL_OVERRIDE, ), StreamParityTestCase(plain=GEMINI_GEMINI_3_8_FLASH_TOOL_CALL, streamed=GEMINI_GEMINI_3_8_FLASH_STREAM_TOOL_CALL), + StreamParityTestCase(plain=GEMINI_GEMINI_4_ARGON_INPUT_TEXT, streamed=GEMINI_GEMINI_4_ARGON_STREAM), + StreamParityTestCase( + plain=GEMINI_GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT, streamed=GEMINI_GEMINI_3_1_PRO_PREVIEW_STREAM + ), ) BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( @@ -3638,4 +3723,6 @@ BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( ), ("/v1/responses", "/v1/messages"), ), + *bridge_pairs((GEMINI_GEMINI_4_ARGON_INPUT_TEXT,), ("/v1/responses",)), + *bridge_pairs((GEMINI_GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT,), ("/v1/responses",)), ) diff --git a/tests/integration/cost_calculation/cases/misc.py b/tests/integration/cost_calculation/cases/misc.py index c3dc0225d0d..c9729353c8e 100644 --- a/tests/integration/cost_calculation/cases/misc.py +++ b/tests/integration/cost_calculation/cases/misc.py @@ -2,7 +2,9 @@ PARITY pairs the streamed case with its plain twin where both bill the same row.""" -from typing import Final +from typing import Final, cast + +from pydantic import JsonValue from integration.cost_calculation.cost_tracking_case import ( ALWAYS_OPEN_OFF_PEAK_WINDOW, @@ -20,13 +22,91 @@ from integration.cost_calculation.cost_tracking_case import ( ) from integration.cost_calculation.stream_parity.case import StreamParityTestCase, sse_frames from integration.cost_calculation.cases.anthropic import ( + CLAUDE_FABLE_5_1_CACHE_READ, + CLAUDE_FABLE_5_1_CACHE_WRITE_1H, + CLAUDE_FABLE_5_1_CACHE_WRITE_5M, CLAUDE_FABLE_5_1_INPUT_TEXT, CLAUDE_FABLE_5_1_STREAM, + CLAUDE_FABLE_5_1_THINKING, + CLAUDE_FABLE_5_1_TOOL_CALL, + CLAUDE_HAIKU_5_5_CACHE_READ, + CLAUDE_HAIKU_5_5_CACHE_WRITE_1H, + CLAUDE_HAIKU_5_5_CACHE_WRITE_5M, CLAUDE_HAIKU_5_5_INPUT_TEXT, CLAUDE_HAIKU_5_5_STREAM, + CLAUDE_HAIKU_5_5_THINKING, + CLAUDE_HAIKU_5_5_TOOL_CALL, ) +def _azure_ai_claude_usage( + input_tokens: int, + output_tokens: int, + inference_geo: str, + *, + cache_creation_input_tokens: int = 0, + cache_read_input_tokens: int = 0, + ephemeral_5m_input_tokens: int = 0, + ephemeral_1h_input_tokens: int = 0, + thinking_tokens: int = 0, +) -> dict[str, JsonValue]: + return { + "input_tokens": input_tokens, + "cache_creation_input_tokens": cache_creation_input_tokens, + "cache_read_input_tokens": cache_read_input_tokens, + "cache_creation": { + "ephemeral_5m_input_tokens": ephemeral_5m_input_tokens, + "ephemeral_1h_input_tokens": ephemeral_1h_input_tokens, + }, + "output_tokens": output_tokens, + "output_tokens_details": {"thinking_tokens": thinking_tokens}, + "service_tier": "standard", + "inference_geo": inference_geo, + "speed": "standard", + } + + +def _azure_ai_claude_shape_case( + sibling: CostTrackingTestCase, + name: str, + model: str, + usage: dict[str, JsonValue], + expected: ExactExpected, + response_content: list[dict[str, JsonValue]] | None = None, + adaptive_thinking: bool = False, +) -> CostTrackingTestCase: + response: Final = cast(JsonResponse, sibling.response) + request: Final = ( + { + **{ + key: value + for key, value in sibling.request.items() + if key not in {"reasoning_effort", "allowed_openai_params"} + }, + "thinking": {"type": "adaptive"}, + "output_config": {"effort": "medium"}, + "allowed_openai_params": ["thinking", "output_config"], + } + if adaptive_thinking + else sibling.request + ) + body: Final = { + **response.body, + "model": model.rsplit("/", 1)[-1], + "usage": usage, + **({"content": response_content} if response_content is not None else {}), + } + return sibling.model_copy( + update={ + "name": name, + "model": model, + "request": request, + "response": response.model_copy(update={"body": body}), + "expected": expected, + } + ) + + # nova-next NOVA_NEXT_TRANSCRIPTIONS_PER_SECOND: Final = CostTrackingTestCase( name="nova-next-transcriptions-per-second", @@ -1457,6 +1537,7 @@ AZURE_AI_CLAUDE_FABLE_5_1_INPUT_TEXT: Final = CLAUDE_FABLE_5_1_INPUT_TEXT.model_ "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "not_available", + "speed": "standard", }, }, ), @@ -1495,6 +1576,7 @@ AZURE_AI_CLAUDE_FABLE_5_1_STREAM: Final = CLAUDE_FABLE_5_1_STREAM.model_copy( "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, "service_tier": "standard", "inference_geo": "not_available", + "speed": "standard", }, }, }, @@ -1582,6 +1664,7 @@ AZURE_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT: Final = CLAUDE_HAIKU_5_5_INPUT_TEXT.model_ "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -1624,6 +1707,7 @@ AZURE_AI_CLAUDE_HAIKU_5_5_STREAM: Final = CLAUDE_HAIKU_5_5_STREAM.model_copy( }, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, }, }, @@ -1713,6 +1797,157 @@ PERPLEXITY_SONAR_NEXT_OFF_PEAK: Final = PERPLEXITY_SONAR_NEXT_NO_SEARCH.model_co } ) + +AZURE_AI_CLAUDE_FABLE_5_1_CACHE_READ: Final = _azure_ai_claude_shape_case( + CLAUDE_FABLE_5_1_CACHE_READ, + "azure_ai/claude-fable-5-1-cache_read", + "azure_ai/claude-fable-5-1", + _azure_ai_claude_usage(input_tokens=14, output_tokens=6, inference_geo="not_available", cache_read_input_tokens=27003), + ExactExpected( + spend=0.012224275, + input_cost=0.011714275, + output_cost=0.00051, + prompt_tokens=27017, + completion_tokens=6, + ), +) +AZURE_AI_CLAUDE_FABLE_5_1_CACHE_WRITE_5M: Final = _azure_ai_claude_shape_case( + CLAUDE_FABLE_5_1_CACHE_WRITE_5M, + "azure_ai/claude-fable-5-1-cache_write_5m", + "azure_ai/claude-fable-5-1", + _azure_ai_claude_usage( + input_tokens=14, + output_tokens=5, + inference_geo="not_available", + cache_creation_input_tokens=27011, + ephemeral_5m_input_tokens=27011, + ), + ExactExpected( + spend=0.57464675, + input_cost=0.57422175, + output_cost=0.000425, + prompt_tokens=27025, + completion_tokens=5, + ), +) +AZURE_AI_CLAUDE_FABLE_5_1_CACHE_WRITE_1H: Final = _azure_ai_claude_shape_case( + CLAUDE_FABLE_5_1_CACHE_WRITE_1H, + "azure_ai/claude-fable-5-1-cache_write_1h", + "azure_ai/claude-fable-5-1", + _azure_ai_claude_usage( + input_tokens=14, + output_tokens=5, + inference_geo="not_available", + cache_creation_input_tokens=27011, + ephemeral_1h_input_tokens=27011, + ), + ExactExpected( + spend=0.919037, + input_cost=0.918612, + output_cost=0.000425, + prompt_tokens=27025, + completion_tokens=5, + ), +) +AZURE_AI_CLAUDE_FABLE_5_1_THINKING: Final = _azure_ai_claude_shape_case( + CLAUDE_FABLE_5_1_THINKING, + "azure_ai/claude-fable-5-1-thinking", + "azure_ai/claude-fable-5-1", + _azure_ai_claude_usage(input_tokens=18, output_tokens=99, inference_geo="not_available", thinking_tokens=33), + ExactExpected( + spend=0.008721, input_cost=0.000306, output_cost=0.008415, prompt_tokens=18, completion_tokens=99 + ), + [ + {"type": "thinking", "thinking": "scripted thought", "signature": "scripted signature"}, + {"type": "text", "text": "scripted answer for cost fixtures"}, + ], + adaptive_thinking=True, +) +AZURE_AI_CLAUDE_FABLE_5_1_TOOL_CALL: Final = _azure_ai_claude_shape_case( + CLAUDE_FABLE_5_1_TOOL_CALL, + "azure_ai/claude-fable-5-1-tool_call", + "azure_ai/claude-fable-5-1", + _azure_ai_claude_usage(input_tokens=385, output_tokens=50, inference_geo="not_available"), + ExactExpected( + spend=0.010795, input_cost=0.006545, output_cost=0.00425, prompt_tokens=385, completion_tokens=50 + ), +) +AZURE_AI_CLAUDE_HAIKU_5_5_CACHE_READ: Final = _azure_ai_claude_shape_case( + CLAUDE_HAIKU_5_5_CACHE_READ, + "azure_ai/claude-haiku-5-5-cache_read", + "azure_ai/claude-haiku-5-5", + _azure_ai_claude_usage(input_tokens=14, output_tokens=5, inference_geo="global", cache_read_input_tokens=27003), + ExactExpected( + spend=0.000520467, + input_cost=0.000515717, + output_cost=0.00000475, + prompt_tokens=27017, + completion_tokens=5, + ), +) +AZURE_AI_CLAUDE_HAIKU_5_5_CACHE_WRITE_5M: Final = _azure_ai_claude_shape_case( + CLAUDE_HAIKU_5_5_CACHE_WRITE_5M, + "azure_ai/claude-haiku-5-5-cache_write_5m", + "azure_ai/claude-haiku-5-5", + _azure_ai_claude_usage( + input_tokens=14, + output_tokens=6, + inference_geo="global", + cache_creation_input_tokens=27011, + ephemeral_5m_input_tokens=27011, + ), + ExactExpected( + spend=0.0064234725, + input_cost=0.0064177725, + output_cost=0.0000057, + prompt_tokens=27025, + completion_tokens=6, + ), +) +AZURE_AI_CLAUDE_HAIKU_5_5_CACHE_WRITE_1H: Final = _azure_ai_claude_shape_case( + CLAUDE_HAIKU_5_5_CACHE_WRITE_1H, + "azure_ai/claude-haiku-5-5-cache_write_1h", + "azure_ai/claude-haiku-5-5", + _azure_ai_claude_usage( + input_tokens=14, + output_tokens=6, + inference_geo="global", + cache_creation_input_tokens=27011, + ephemeral_1h_input_tokens=27011, + ), + ExactExpected( + spend=0.01027254, + input_cost=0.01026684, + output_cost=0.0000057, + prompt_tokens=27025, + completion_tokens=6, + ), +) +AZURE_AI_CLAUDE_HAIKU_5_5_THINKING: Final = _azure_ai_claude_shape_case( + CLAUDE_HAIKU_5_5_THINKING, + "azure_ai/claude-haiku-5-5-thinking", + "azure_ai/claude-haiku-5-5", + _azure_ai_claude_usage(input_tokens=18, output_tokens=95, inference_geo="global", thinking_tokens=54), + ExactExpected( + spend=0.00009367, input_cost=0.00000342, output_cost=0.00009025, prompt_tokens=18, completion_tokens=95 + ), + [ + {"type": "thinking", "thinking": "scripted thought", "signature": "scripted signature"}, + {"type": "text", "text": "scripted answer for cost fixtures"}, + ], + adaptive_thinking=True, +) +AZURE_AI_CLAUDE_HAIKU_5_5_TOOL_CALL: Final = _azure_ai_claude_shape_case( + CLAUDE_HAIKU_5_5_TOOL_CALL, + "azure_ai/claude-haiku-5-5-tool_call", + "azure_ai/claude-haiku-5-5", + _azure_ai_claude_usage(input_tokens=511, output_tokens=34, inference_geo="global"), + ExactExpected( + spend=0.00012939, input_cost=0.00009709, output_cost=0.0000323, prompt_tokens=511, completion_tokens=34 + ), +) + + CASES: Final[tuple[CostTrackingTestCase, ...]] = ( PERPLEXITY_SONAR_NEXT_OFF_PEAK, NOVA_NEXT_TRANSCRIPTIONS_PER_SECOND, @@ -1764,8 +1999,18 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( XAI_GROK_4_7_INPUT_TEXT, AZURE_AI_CLAUDE_FABLE_5_1_INPUT_TEXT, AZURE_AI_CLAUDE_FABLE_5_1_STREAM, + AZURE_AI_CLAUDE_FABLE_5_1_CACHE_READ, + AZURE_AI_CLAUDE_FABLE_5_1_CACHE_WRITE_5M, + AZURE_AI_CLAUDE_FABLE_5_1_CACHE_WRITE_1H, + AZURE_AI_CLAUDE_FABLE_5_1_THINKING, + AZURE_AI_CLAUDE_FABLE_5_1_TOOL_CALL, AZURE_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT, AZURE_AI_CLAUDE_HAIKU_5_5_STREAM, + AZURE_AI_CLAUDE_HAIKU_5_5_CACHE_READ, + AZURE_AI_CLAUDE_HAIKU_5_5_CACHE_WRITE_5M, + AZURE_AI_CLAUDE_HAIKU_5_5_CACHE_WRITE_1H, + AZURE_AI_CLAUDE_HAIKU_5_5_THINKING, + AZURE_AI_CLAUDE_HAIKU_5_5_TOOL_CALL, ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( diff --git a/tests/integration/cost_calculation/cases/openai.py b/tests/integration/cost_calculation/cases/openai.py index cfb7fda4d71..85ac05dd69d 100644 --- a/tests/integration/cost_calculation/cases/openai.py +++ b/tests/integration/cost_calculation/cases/openai.py @@ -26,6 +26,14 @@ from integration.cost_calculation.stream_parity.case import StreamParityTestCase from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParityTestCase, bridge_pairs +def _response_for_model(response: JsonResponse, model: str) -> JsonResponse: + return response.model_copy(update={"body": {**response.body, "model": model}}) + + +def _stream_for_model(response: SseResponse, model: str) -> SseResponse: + return response.model_copy(update={"frames": tuple(frame.replace("gpt-6-luna", model) for frame in response.frames)}) + + # gpt-5.3-codex GPT_5_3_CODEX_INPUT_TEXT: Final = CostTrackingTestCase( name="gpt-5.3-codex-input_text", @@ -11183,6 +11191,171 @@ GPT_6_ASTRA_CACHE_WRITE_5M: Final = with_response_usage( ExactExpected(spend=0.2148, input_cost=0.192, output_cost=0.0228, prompt_tokens=12928, completion_tokens=380), ) +GPT_5_6_TERRA_INPUT_TEXT: Final = GPT_6_LUNA_INPUT_TEXT.model_copy( + update={ + "name": "gpt-5.6-terra-input_text", + "model": "gpt-5.6-terra", + "response": _response_for_model(cast(JsonResponse, GPT_6_LUNA_INPUT_TEXT.response), "gpt-5.6-terra"), + "expected": ExactExpected( + spend=0.0112112, + input_cost=0.004784, + output_cost=0.0064272, + prompt_tokens=1840, + completion_tokens=412, + ), + } +) +GPT_5_6_TERRA_STREAM: Final = GPT_6_LUNA_STREAM.model_copy( + update={ + "name": "gpt-5.6-terra-stream", + "model": "gpt-5.6-terra", + "response": _stream_for_model(cast(SseResponse, GPT_6_LUNA_STREAM.response), "gpt-5.6-terra"), + "expected": ExactExpected( + spend=0.0112112, + input_cost=0.004784, + output_cost=0.0064272, + prompt_tokens=1840, + completion_tokens=412, + ), + } +) +GPT_5_6_TERRA_SERVICE_TIER_FLEX: Final = GPT_6_LUNA_SERVICE_TIER_FLEX.model_copy( + update={ + "name": "gpt-5.6-terra-service_tier_flex", + "model": "gpt-5.6-terra", + "response": _response_for_model(cast(JsonResponse, GPT_6_LUNA_SERVICE_TIER_FLEX.response), "gpt-5.6-terra"), + "expected": ExactExpected( + spend=0.0056056, + input_cost=0.002392, + output_cost=0.0032136, + prompt_tokens=1840, + completion_tokens=412, + ), + } +) +GPT_5_6_TERRA_SERVICE_TIER_PRIORITY: Final = GPT_6_LUNA_SERVICE_TIER_PRIORITY.model_copy( + update={ + "name": "gpt-5.6-terra-service_tier_priority", + "model": "gpt-5.6-terra", + "response": _response_for_model( + cast(JsonResponse, GPT_6_LUNA_SERVICE_TIER_PRIORITY.response), "gpt-5.6-terra" + ), + "expected": ExactExpected( + spend=0.0224224, + input_cost=0.009568, + output_cost=0.0128544, + prompt_tokens=1840, + completion_tokens=412, + ), + } +) +GPT_5_6_TERRA_RESPONSES_NATIVE_JSON: Final = GPT_6_LUNA_RESPONSES_NATIVE_JSON.model_copy( + update={ + "name": "gpt-5.6-terra-responses_native_json", + "model": "gpt-5.6-terra", + "response": _response_for_model( + cast(JsonResponse, GPT_6_LUNA_RESPONSES_NATIVE_JSON.response), "gpt-5.6-terra" + ), + "expected": ExactExpected( + spend=0.0112112, + input_cost=0.004784, + output_cost=0.0064272, + prompt_tokens=1840, + completion_tokens=412, + ), + } +) +GPT_5_6_TERRA_RESPONSES_DEPLOYMENT_JSON: Final = GPT_6_LUNA_RESPONSES_DEPLOYMENT_JSON.model_copy( + update={ + "name": "gpt-5.6-terra-responses_deployment_json", + "model": "gpt-5.6-terra", + "deployment": Deployment(model="openai/responses/gpt-5.6-terra"), + "response": _response_for_model( + cast(JsonResponse, GPT_6_LUNA_RESPONSES_DEPLOYMENT_JSON.response), "gpt-5.6-terra" + ), + "expected": ExactExpected( + spend=0.0112112, + input_cost=0.004784, + output_cost=0.0064272, + prompt_tokens=1840, + completion_tokens=412, + ), + } +) +GPT_5_6_TERRA_CHAT_BRIDGED_JSON: Final = GPT_6_LUNA_CHAT_BRIDGED_JSON.model_copy( + update={ + "name": "gpt-5.6-terra-chat_bridged_json", + "model": "gpt-5.6-terra", + "deployment": Deployment(model="openai/responses/gpt-5.6-terra"), + "response": _response_for_model(cast(JsonResponse, GPT_6_LUNA_CHAT_BRIDGED_JSON.response), "gpt-5.6-terra"), + "expected": ExactExpected( + spend=0.0112112, + input_cost=0.004784, + output_cost=0.0064272, + prompt_tokens=1840, + completion_tokens=412, + ), + } +) +GPT_5_6_TERRA_CACHE_READ: Final = GPT_6_LUNA_CACHE_READ.model_copy( + update={ + "name": "gpt-5.6-terra-cache_read", + "model": "gpt-5.6-terra", + "response": _response_for_model(cast(JsonResponse, GPT_6_LUNA_CACHE_READ.response), "gpt-5.6-terra"), + "expected": ExactExpected( + spend=0.01078688, + input_cost=0.00485888, + output_cost=0.005928, + prompt_tokens=12928, + completion_tokens=380, + ), + } +) +GPT_5_6_TERRA_CACHE_WRITE_5M: Final = GPT_6_LUNA_CACHE_WRITE_5M.model_copy( + update={ + "name": "gpt-5.6-terra-cache_write_5m", + "model": "gpt-5.6-terra", + "response": _response_for_model( + cast(JsonResponse, GPT_6_LUNA_CACHE_WRITE_5M.response), "gpt-5.6-terra" + ), + "expected": ExactExpected( + spend=0.047528, + input_cost=0.0416, + output_cost=0.005928, + prompt_tokens=12928, + completion_tokens=380, + ), + } +) +GPT_5_6_TERRA_REASONING: Final = GPT_6_LUNA_REASONING.model_copy( + update={ + "name": "gpt-5.6-terra-reasoning", + "model": "gpt-5.6-terra", + "response": _response_for_model(cast(JsonResponse, GPT_6_LUNA_REASONING.response), "gpt-5.6-terra"), + "expected": ExactExpected( + spend=0.066248, + input_cost=0.003224, + output_cost=0.063024, + prompt_tokens=1240, + completion_tokens=4040, + ), + } +) +GPT_5_6_TERRA_TOOL_CALL: Final = GPT_6_LUNA_TOOL_CALL.model_copy( + update={ + "name": "gpt-5.6-terra-tool_call", + "model": "gpt-5.6-terra", + "response": _response_for_model(cast(JsonResponse, GPT_6_LUNA_TOOL_CALL.response), "gpt-5.6-terra"), + "expected": ExactExpected( + spend=0.0112112, + input_cost=0.004784, + output_cost=0.0064272, + prompt_tokens=1840, + completion_tokens=412, + ), + } +) + GPT_5_6_SLASH_ALIAS: Final = GPT_5_6_INPUT_TEXT.model_copy( update={"name": "gpt-5.6-slash_alias", "alias": "openai/custom/gpt-5.6-$MARKER"} @@ -11485,6 +11658,17 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GPT_6_1_SOL_CACHE_WRITE_5M, GPT_6_LUNA_CACHE_WRITE_5M, GPT_6_ASTRA_CACHE_WRITE_5M, + GPT_5_6_TERRA_INPUT_TEXT, + GPT_5_6_TERRA_STREAM, + GPT_5_6_TERRA_RESPONSES_NATIVE_JSON, + GPT_5_6_TERRA_RESPONSES_DEPLOYMENT_JSON, + GPT_5_6_TERRA_CHAT_BRIDGED_JSON, + GPT_5_6_TERRA_CACHE_READ, + GPT_5_6_TERRA_CACHE_WRITE_5M, + GPT_5_6_TERRA_REASONING, + GPT_5_6_TERRA_SERVICE_TIER_FLEX, + GPT_5_6_TERRA_SERVICE_TIER_PRIORITY, + GPT_5_6_TERRA_TOOL_CALL, ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( @@ -11513,6 +11697,7 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( StreamParityTestCase(plain=GPT_6_LUNA_INPUT_TEXT, streamed=GPT_6_LUNA_STREAM), StreamParityTestCase(plain=GPT_6_ASTRA_INPUT_TEXT, streamed=GPT_6_ASTRA_STREAM), StreamParityTestCase(plain=GPT_6_1_SOL_INPUT_TEXT, streamed=GPT_6_1_SOL_STREAM), + StreamParityTestCase(plain=GPT_5_6_TERRA_INPUT_TEXT, streamed=GPT_5_6_TERRA_STREAM), ) BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( @@ -11545,4 +11730,7 @@ BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( BridgeCostParityTestCase(native=GPT_6_LUNA_RESPONSES_DEPLOYMENT_JSON, bridged=GPT_6_LUNA_CHAT_BRIDGED_JSON), BridgeCostParityTestCase(native=GPT_6_ASTRA_RESPONSES_DEPLOYMENT_JSON, bridged=GPT_6_ASTRA_CHAT_BRIDGED_JSON), BridgeCostParityTestCase(native=GPT_6_1_SOL_RESPONSES_DEPLOYMENT_JSON, bridged=GPT_6_1_SOL_CHAT_BRIDGED_JSON), + BridgeCostParityTestCase( + native=GPT_5_6_TERRA_RESPONSES_DEPLOYMENT_JSON, bridged=GPT_5_6_TERRA_CHAT_BRIDGED_JSON + ), ) diff --git a/tests/integration/cost_calculation/cases/vertex_ai.py b/tests/integration/cost_calculation/cases/vertex_ai.py index 966a24bcba2..e35fd2307c2 100644 --- a/tests/integration/cost_calculation/cases/vertex_ai.py +++ b/tests/integration/cost_calculation/cases/vertex_ai.py @@ -2,15 +2,32 @@ PARITY pairs the streamed case with its plain twin where both bill the same row.""" -from typing import Final +from typing import Final, cast + +from pydantic import JsonValue from integration.cost_calculation.cases.anthropic import ( + CLAUDE_FABLE_5_1_CACHE_READ, + CLAUDE_FABLE_5_1_CACHE_WRITE_1H, + CLAUDE_FABLE_5_1_CACHE_WRITE_5M, CLAUDE_FABLE_5_1_INPUT_TEXT, CLAUDE_FABLE_5_1_STREAM, + CLAUDE_FABLE_5_1_THINKING, + CLAUDE_FABLE_5_1_TOOL_CALL, + CLAUDE_OPUS_5_5_CACHE_READ, + CLAUDE_OPUS_5_5_CACHE_WRITE_1H, + CLAUDE_OPUS_5_5_CACHE_WRITE_5M, CLAUDE_OPUS_5_5_INPUT_TEXT, CLAUDE_OPUS_5_5_STREAM, + CLAUDE_OPUS_5_5_THINKING, + CLAUDE_OPUS_5_5_TOOL_CALL, + CLAUDE_SONNET_5_5_CACHE_READ, + CLAUDE_SONNET_5_5_CACHE_WRITE_1H, + CLAUDE_SONNET_5_5_CACHE_WRITE_5M, CLAUDE_SONNET_5_5_INPUT_TEXT, CLAUDE_SONNET_5_5_STREAM, + CLAUDE_SONNET_5_5_THINKING, + CLAUDE_SONNET_5_5_TOOL_CALL, ) from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParityTestCase, bridge_pairs from integration.cost_calculation.cost_tracking_case import ( @@ -24,6 +41,99 @@ from integration.cost_calculation.cost_tracking_case import ( from integration.cost_calculation.stream_parity.case import StreamParityTestCase, sse_frames +def _gemini_case_response(response: JsonResponse | SseResponse, model_version: str) -> JsonResponse | SseResponse: + match response: + case JsonResponse(): + return response.model_copy(update={"body": {**response.body, "modelVersion": model_version}}) + case SseResponse(): + return response.model_copy( + update={ + "frames": tuple(frame.replace("gemini-3.8-flash", model_version) for frame in response.frames) + } + ) + + +def _gemini_case( + sibling: CostTrackingTestCase, name: str, model: str, expected: ExactExpected +) -> CostTrackingTestCase: + return sibling.model_copy( + update={ + "name": name, + "model": model, + "response": _gemini_case_response(cast(JsonResponse | SseResponse, sibling.response), "gemini-3.1-pro-preview"), + "expected": expected, + } + ) + + +def _vertex_claude_usage( + input_tokens: int, + output_tokens: int, + *, + cache_creation_input_tokens: int = 0, + cache_read_input_tokens: int = 0, + ephemeral_5m_input_tokens: int = 0, + ephemeral_1h_input_tokens: int = 0, + thinking_tokens: int = 0, +) -> dict[str, JsonValue]: + return { + "input_tokens": input_tokens, + "cache_creation_input_tokens": cache_creation_input_tokens, + "cache_read_input_tokens": cache_read_input_tokens, + "cache_creation": { + "ephemeral_5m_input_tokens": ephemeral_5m_input_tokens, + "ephemeral_1h_input_tokens": ephemeral_1h_input_tokens, + }, + "output_tokens": output_tokens, + "output_tokens_details": {"thinking_tokens": thinking_tokens}, + } + + +def _vertex_claude_shape_case( + sibling: CostTrackingTestCase, + name: str, + model: str, + usage: dict[str, JsonValue], + expected: ExactExpected, + response_content: list[dict[str, JsonValue]] | None = None, + adaptive_thinking: bool = False, +) -> CostTrackingTestCase: + response: Final = cast(JsonResponse, sibling.response) + request: Final = ( + { + **{ + key: value + for key, value in sibling.request.items() + if key not in {"reasoning_effort", "allowed_openai_params"} + }, + "thinking": {"type": "adaptive"}, + "output_config": {"effort": "medium"}, + "allowed_openai_params": ["thinking", "output_config"], + } + if adaptive_thinking + else sibling.request + ) + body: Final = { + key: value + for key, value in { + **response.body, + "model": model.rsplit("/", 1)[-1], + "usage": usage, + **({"content": response_content} if response_content is not None else {}), + }.items() + if key != "diagnostics" + } + return sibling.model_copy( + update={ + "name": name, + "model": model, + "request": request, + "response": response.model_copy(update={"body": body}), + "expected": expected, + } + ) + + def _vertex_claude_stream_response(model: str, stop_reason: str, text_deltas: tuple[str, ...]) -> SseResponse: text_frames: Final = ( ( @@ -3081,6 +3191,191 @@ VERTEX_AI_CLAUDE_FABLE_5_1_STREAM: Final = CLAUDE_FABLE_5_1_STREAM.model_copy( } ) +VERTEX_AI_CLAUDE_SONNET_5_5_CACHE_READ: Final = _vertex_claude_shape_case( + CLAUDE_SONNET_5_5_CACHE_READ, + "vertex_ai/claude-sonnet-5-5-cache_read", + "vertex_ai/claude-sonnet-5-5", + _vertex_claude_usage(input_tokens=14, output_tokens=6, cache_read_input_tokens=27003), + ExactExpected( + spend=0.00474011, input_cost=0.00463811, output_cost=0.000102, prompt_tokens=27017, completion_tokens=6 + ), +) +VERTEX_AI_CLAUDE_SONNET_5_5_CACHE_WRITE_5M: Final = _vertex_claude_shape_case( + CLAUDE_SONNET_5_5_CACHE_WRITE_5M, + "vertex_ai/claude-sonnet-5-5-cache_write_5m", + "vertex_ai/claude-sonnet-5-5", + _vertex_claude_usage( + input_tokens=14, output_tokens=6, cache_creation_input_tokens=27011, ephemeral_5m_input_tokens=27011 + ), + ExactExpected( + spend=0.11494635, input_cost=0.11484435, output_cost=0.000102, prompt_tokens=27025, completion_tokens=6 + ), +) +VERTEX_AI_CLAUDE_SONNET_5_5_CACHE_WRITE_1H: Final = _vertex_claude_shape_case( + CLAUDE_SONNET_5_5_CACHE_WRITE_1H, + "vertex_ai/claude-sonnet-5-5-cache_write_1h", + "vertex_ai/claude-sonnet-5-5", + _vertex_claude_usage( + input_tokens=14, output_tokens=6, cache_creation_input_tokens=27011, ephemeral_1h_input_tokens=27011 + ), + ExactExpected( + spend=0.1838244, input_cost=0.1837224, output_cost=0.000102, prompt_tokens=27025, completion_tokens=6 + ), +) +VERTEX_AI_CLAUDE_SONNET_5_5_THINKING: Final = _vertex_claude_shape_case( + CLAUDE_SONNET_5_5_THINKING, + "vertex_ai/claude-sonnet-5-5-thinking", + "vertex_ai/claude-sonnet-5-5", + _vertex_claude_usage(input_tokens=18, output_tokens=144, thinking_tokens=37), + ExactExpected( + spend=0.0025092, input_cost=0.0000612, output_cost=0.002448, prompt_tokens=18, completion_tokens=144 + ), + [ + {"type": "thinking", "thinking": "scripted thought", "signature": "scripted signature"}, + {"type": "text", "text": "scripted answer for cost fixtures"}, + ], + adaptive_thinking=True, +) +VERTEX_AI_CLAUDE_SONNET_5_5_TOOL_CALL: Final = _vertex_claude_shape_case( + CLAUDE_SONNET_5_5_TOOL_CALL, + "vertex_ai/claude-sonnet-5-5-tool_call", + "vertex_ai/claude-sonnet-5-5", + _vertex_claude_usage(input_tokens=385, output_tokens=50), + ExactExpected( + spend=0.002159, input_cost=0.001309, output_cost=0.00085, prompt_tokens=385, completion_tokens=50 + ), +) +VERTEX_AI_CLAUDE_OPUS_5_5_CACHE_READ: Final = _vertex_claude_shape_case( + CLAUDE_OPUS_5_5_CACHE_READ, + "vertex_ai/claude-opus-5-5-cache_read", + "vertex_ai/claude-opus-5-5", + _vertex_claude_usage(input_tokens=14, output_tokens=6, cache_read_input_tokens=27003), + ExactExpected( + spend=0.01003788, input_cost=0.00982188, output_cost=0.000216, prompt_tokens=27017, completion_tokens=6 + ), +) +VERTEX_AI_CLAUDE_OPUS_5_5_CACHE_WRITE_5M: Final = _vertex_claude_shape_case( + CLAUDE_OPUS_5_5_CACHE_WRITE_5M, + "vertex_ai/claude-opus-5-5-cache_write_5m", + "vertex_ai/claude-opus-5-5", + _vertex_claude_usage( + input_tokens=14, output_tokens=6, cache_creation_input_tokens=27011, ephemeral_5m_input_tokens=27011 + ), + ExactExpected( + spend=0.2434158, input_cost=0.2431998, output_cost=0.000216, prompt_tokens=27025, completion_tokens=6 + ), +) +VERTEX_AI_CLAUDE_OPUS_5_5_CACHE_WRITE_1H: Final = _vertex_claude_shape_case( + CLAUDE_OPUS_5_5_CACHE_WRITE_1H, + "vertex_ai/claude-opus-5-5-cache_write_1h", + "vertex_ai/claude-opus-5-5", + _vertex_claude_usage( + input_tokens=14, output_tokens=6, cache_creation_input_tokens=27011, ephemeral_1h_input_tokens=27011 + ), + ExactExpected( + spend=0.3892752, input_cost=0.3890592, output_cost=0.000216, prompt_tokens=27025, completion_tokens=6 + ), +) +VERTEX_AI_CLAUDE_OPUS_5_5_THINKING: Final = _vertex_claude_shape_case( + CLAUDE_OPUS_5_5_THINKING, + "vertex_ai/claude-opus-5-5-thinking", + "vertex_ai/claude-opus-5-5", + _vertex_claude_usage(input_tokens=18, output_tokens=162, thinking_tokens=66), + ExactExpected( + spend=0.0059616, input_cost=0.0001296, output_cost=0.005832, prompt_tokens=18, completion_tokens=162 + ), + [ + {"type": "thinking", "thinking": "scripted thought", "signature": "scripted signature"}, + {"type": "text", "text": "scripted answer for cost fixtures"}, + ], + adaptive_thinking=True, +) +VERTEX_AI_CLAUDE_OPUS_5_5_TOOL_CALL: Final = _vertex_claude_shape_case( + CLAUDE_OPUS_5_5_TOOL_CALL, + "vertex_ai/claude-opus-5-5-tool_call", + "vertex_ai/claude-opus-5-5", + _vertex_claude_usage(input_tokens=385, output_tokens=50), + ExactExpected( + spend=0.004572, input_cost=0.002772, output_cost=0.0018, prompt_tokens=385, completion_tokens=50 + ), +) +VERTEX_AI_CLAUDE_FABLE_5_1_CACHE_READ: Final = _vertex_claude_shape_case( + CLAUDE_FABLE_5_1_CACHE_READ, + "vertex_ai/claude-fable-5-1-cache_read", + "vertex_ai/claude-fable-5-1", + _vertex_claude_usage(input_tokens=14, output_tokens=5, cache_read_input_tokens=27003), + ExactExpected( + spend=0.013567425, input_cost=0.013092425, output_cost=0.000475, prompt_tokens=27017, completion_tokens=5 + ), +) +VERTEX_AI_CLAUDE_FABLE_5_1_CACHE_WRITE_5M: Final = _vertex_claude_shape_case( + CLAUDE_FABLE_5_1_CACHE_WRITE_5M, + "vertex_ai/claude-fable-5-1-cache_write_5m", + "vertex_ai/claude-fable-5-1", + _vertex_claude_usage( + input_tokens=14, output_tokens=5, cache_creation_input_tokens=27011, ephemeral_5m_input_tokens=27011 + ), + ExactExpected( + spend=0.64225225, input_cost=0.64177725, output_cost=0.000475, prompt_tokens=27025, completion_tokens=5 + ), +) +VERTEX_AI_CLAUDE_FABLE_5_1_CACHE_WRITE_1H: Final = _vertex_claude_shape_case( + CLAUDE_FABLE_5_1_CACHE_WRITE_1H, + "vertex_ai/claude-fable-5-1-cache_write_1h", + "vertex_ai/claude-fable-5-1", + _vertex_claude_usage( + input_tokens=14, output_tokens=5, cache_creation_input_tokens=27011, ephemeral_1h_input_tokens=27011 + ), + ExactExpected( + spend=1.027159, input_cost=1.026684, output_cost=0.000475, prompt_tokens=27025, completion_tokens=5 + ), +) +VERTEX_AI_CLAUDE_FABLE_5_1_THINKING: Final = _vertex_claude_shape_case( + CLAUDE_FABLE_5_1_THINKING, + "vertex_ai/claude-fable-5-1-thinking", + "vertex_ai/claude-fable-5-1", + _vertex_claude_usage(input_tokens=18, output_tokens=98, thinking_tokens=33), + ExactExpected( + spend=0.009652, input_cost=0.000342, output_cost=0.00931, prompt_tokens=18, completion_tokens=98 + ), + [ + {"type": "thinking", "thinking": "scripted thought", "signature": "scripted signature"}, + {"type": "text", "text": "scripted answer for cost fixtures"}, + ], + adaptive_thinking=True, +) +VERTEX_AI_CLAUDE_FABLE_5_1_TOOL_CALL: Final = _vertex_claude_shape_case( + CLAUDE_FABLE_5_1_TOOL_CALL, + "vertex_ai/claude-fable-5-1-tool_call", + "vertex_ai/claude-fable-5-1", + _vertex_claude_usage(input_tokens=385, output_tokens=50), + ExactExpected( + spend=0.012065, input_cost=0.007315, output_cost=0.00475, prompt_tokens=385, completion_tokens=50 + ), +) + + +GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT: Final = _gemini_case( + GEMINI_3_8_FLASH_INPUT_TEXT, "gemini-3.1-pro-preview-input_text", "gemini-3.1-pro-preview", + ExactExpected(spend=0.0094864, input_cost=0.004048, output_cost=0.0054384, prompt_tokens=1840, completion_tokens=412) +) +GEMINI_3_1_PRO_PREVIEW_STREAM: Final = _gemini_case( + GEMINI_3_8_FLASH_STREAM, "gemini-3.1-pro-preview-stream", "gemini-3.1-pro-preview", + ExactExpected(spend=0.0094864, input_cost=0.004048, output_cost=0.0054384, prompt_tokens=1840, completion_tokens=412) +) +GEMINI_3_1_PRO_PREVIEW_CACHE_READ: Final = _gemini_case( + GEMINI_3_8_FLASH_CACHE_READ, "gemini-3.1-pro-preview-cache_read", "gemini-3.1-pro-preview", + ExactExpected(spend=0.00912736, input_cost=0.00411136, output_cost=0.005016, prompt_tokens=12928, completion_tokens=380) +) +GEMINI_3_1_PRO_PREVIEW_REASONING: Final = _gemini_case( + GEMINI_3_8_FLASH_FALLBACK_REASONING_AT_OUTPUT_RATE, "gemini-3.1-pro-preview-reasoning", "gemini-3.1-pro-preview", + ExactExpected(spend=0.056056, input_cost=0.002728, output_cost=0.053328, prompt_tokens=1240, completion_tokens=4040) +) +GEMINI_3_1_PRO_PREVIEW_TOOL_CALL: Final = _gemini_case( + GEMINI_3_8_FLASH_TOOL_CALL, "gemini-3.1-pro-preview-tool_call", "gemini-3.1-pro-preview", + ExactExpected(spend=0.0094864, input_cost=0.004048, output_cost=0.0054384, prompt_tokens=1840, completion_tokens=412) +) + CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GEMINI_3_1_PRO_INPUT_TEXT, @@ -3137,6 +3432,26 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( VERTEX_AI_CLAUDE_OPUS_5_5_STREAM, VERTEX_AI_CLAUDE_FABLE_5_1_INPUT_TEXT, VERTEX_AI_CLAUDE_FABLE_5_1_STREAM, + VERTEX_AI_CLAUDE_SONNET_5_5_CACHE_READ, + VERTEX_AI_CLAUDE_SONNET_5_5_CACHE_WRITE_5M, + VERTEX_AI_CLAUDE_SONNET_5_5_CACHE_WRITE_1H, + VERTEX_AI_CLAUDE_SONNET_5_5_THINKING, + VERTEX_AI_CLAUDE_SONNET_5_5_TOOL_CALL, + VERTEX_AI_CLAUDE_OPUS_5_5_CACHE_READ, + VERTEX_AI_CLAUDE_OPUS_5_5_CACHE_WRITE_5M, + VERTEX_AI_CLAUDE_OPUS_5_5_CACHE_WRITE_1H, + VERTEX_AI_CLAUDE_OPUS_5_5_THINKING, + VERTEX_AI_CLAUDE_OPUS_5_5_TOOL_CALL, + VERTEX_AI_CLAUDE_FABLE_5_1_CACHE_READ, + VERTEX_AI_CLAUDE_FABLE_5_1_CACHE_WRITE_5M, + VERTEX_AI_CLAUDE_FABLE_5_1_CACHE_WRITE_1H, + VERTEX_AI_CLAUDE_FABLE_5_1_THINKING, + VERTEX_AI_CLAUDE_FABLE_5_1_TOOL_CALL, + GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT, + GEMINI_3_1_PRO_PREVIEW_STREAM, + GEMINI_3_1_PRO_PREVIEW_CACHE_READ, + GEMINI_3_1_PRO_PREVIEW_REASONING, + GEMINI_3_1_PRO_PREVIEW_TOOL_CALL, ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( @@ -3157,9 +3472,13 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( ), StreamParityTestCase(plain=VERTEX_AI_CLAUDE_OPUS_5_5_INPUT_TEXT, streamed=VERTEX_AI_CLAUDE_OPUS_5_5_STREAM), StreamParityTestCase(plain=VERTEX_AI_CLAUDE_FABLE_5_1_INPUT_TEXT, streamed=VERTEX_AI_CLAUDE_FABLE_5_1_STREAM), + StreamParityTestCase(plain=GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT, streamed=GEMINI_3_1_PRO_PREVIEW_STREAM), ) -BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = bridge_pairs( - (GEMINI_3_8_FLASH_INPUT_TEXT, GEMINI_3_8_FLASH_CACHE_READ, GEMINI_3_8_FLASH_STREAM), - ("/v1/responses", "/v1/messages"), +BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( + *bridge_pairs( + (GEMINI_3_8_FLASH_INPUT_TEXT, GEMINI_3_8_FLASH_CACHE_READ, GEMINI_3_8_FLASH_STREAM), + ("/v1/responses", "/v1/messages"), + ), + *bridge_pairs((GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT,), ("/v1/responses",)), ) diff --git a/tests/integration/cost_calculation/cost_map.py b/tests/integration/cost_calculation/cost_map.py index 78d025855b5..40072bfcaf6 100644 --- a/tests/integration/cost_calculation/cost_map.py +++ b/tests/integration/cost_calculation/cost_map.py @@ -32,6 +32,41 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "input_cost_per_token": 3.2e-06, "output_cost_per_token": 1.6e-05, }, + "gpt-5.6-terra": { + "cache_creation_input_token_cost": 3.25e-06, + "cache_creation_input_token_cost_above_272k_tokens": 6.5e-06, + "cache_creation_input_token_cost_above_272k_tokens_flex": 3.25e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1.3e-05, + "cache_creation_input_token_cost_flex": 1.625e-06, + "cache_creation_input_token_cost_priority": 6.5e-06, + "cache_read_input_token_cost": 2.6e-07, + "cache_read_input_token_cost_above_272k_tokens": 5.2e-07, + "cache_read_input_token_cost_above_272k_tokens_flex": 2.6e-07, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.04e-06, + "cache_read_input_token_cost_flex": 1.3e-07, + "cache_read_input_token_cost_priority": 5.2e-07, + "input_cost_per_token": 2.6e-06, + "input_cost_per_token_above_272k_tokens": 5.2e-06, + "input_cost_per_token_above_272k_tokens_flex": 2.6e-06, + "input_cost_per_token_above_272k_tokens_priority": 1.04e-05, + "input_cost_per_token_flex": 1.3e-06, + "input_cost_per_token_priority": 5.2e-06, + "litellm_provider": "openai", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.56e-05, + "output_cost_per_token_above_272k_tokens": 2.34e-05, + "output_cost_per_token_above_272k_tokens_flex": 1.17e-05, + "output_cost_per_token_above_272k_tokens_priority": 4.68e-05, + "output_cost_per_token_flex": 7.8e-06, + "output_cost_per_token_priority": 3.12e-05, + "supports_function_calling": True, + "supports_none_reasoning_effort": True, + "supports_prompt_caching": True, + "supports_reasoning": True, + }, "gpt-6-luna": { "output_cost_per_token_priority": 1e-06, "input_cost_per_token_priority": 2e-07, @@ -407,8 +442,55 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "supports_function_calling": True, "web_search_billing_unit": "per_query", }, + "gemini/gemini-4-argon": { + "cache_read_input_token_cost": 1.1e-07, + "cache_read_input_token_cost_flex": 5.5e-08, + "cache_read_input_token_cost_priority": 1.98e-07, + "input_cost_per_token": 2.2e-06, + "input_cost_per_token_flex": 1.1e-06, + "input_cost_per_token_priority": 3.96e-06, + "litellm_provider": "gemini", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_reasoning_token": 1.1e-05, + "output_cost_per_token": 1.1e-05, + "output_cost_per_token_flex": 5.5e-06, + "output_cost_per_token_priority": 1.98e-05, + "supports_function_calling": True, + "supports_prompt_caching": True, + "supports_reasoning": True, + }, + "gemini/gemini-3.1-pro-preview": { + "cache_read_input_token_cost": 2.6e-07, + "cache_read_input_token_cost_above_200k_tokens": 5.2e-07, + "input_cost_per_token": 2.6e-06, + "input_cost_per_token_above_200k_tokens": 5.2e-06, + "litellm_provider": "gemini", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 1.56e-05, + "output_cost_per_token_above_200k_tokens": 2.34e-05, + "supports_function_calling": True, + "supports_prompt_caching": True, + "supports_reasoning": True, + "input_cost_per_token_priority": 4.68e-06, + "input_cost_per_token_above_200k_tokens_priority": 9.36e-06, + "output_cost_per_token_priority": 2.808e-05, + "output_cost_per_token_above_200k_tokens_priority": 4.212e-05, + "cache_read_input_token_cost_priority": 4.68e-07, + "cache_read_input_token_cost_above_200k_tokens_priority": 9.36e-07, + "cache_read_input_token_cost_flex": 2.6e-07, + "input_cost_per_token_flex": 1.3e-06, + "output_cost_per_token_flex": 7.8e-06, + }, "gemini/gemini-3.8-flash": { "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost_flex": 2.5e-08, + "cache_read_input_token_cost_priority": 9e-08, "google_maps_grounding_cost_per_query": 0.025, "input_cost_per_audio_token": 1e-06, "input_cost_per_image_token": 5.5e-07, @@ -454,8 +536,36 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "supports_function_calling": True, "web_search_billing_unit": "per_query", }, + "gemini-3.1-pro-preview": { + "cache_creation_input_token_cost_above_200k_tokens": 2.75e-07, + "cache_read_input_token_cost": 2.2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4.4e-07, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.92e-07, + "cache_read_input_token_cost_priority": 3.96e-07, + "input_cost_per_token": 2.2e-06, + "input_cost_per_token_above_200k_tokens": 4.4e-06, + "input_cost_per_token_above_200k_tokens_priority": 7.92e-06, + "input_cost_per_token_flex": 1.1e-06, + "input_cost_per_token_priority": 3.96e-06, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_image": 0.000132, + "output_cost_per_token": 1.32e-05, + "output_cost_per_token_above_200k_tokens": 1.98e-05, + "output_cost_per_token_above_200k_tokens_priority": 3.564e-05, + "output_cost_per_token_flex": 6.6e-06, + "output_cost_per_token_priority": 2.376e-05, + "supports_function_calling": True, + "supports_prompt_caching": True, + "supports_reasoning": True, + }, "gemini-3.8-flash": { "cache_read_input_token_cost": 5.2e-08, + "cache_read_input_token_cost_flex": 2.6e-08, + "cache_read_input_token_cost_priority": 9.36e-08, "google_maps_grounding_cost_per_query": 0.025, "input_cost_per_audio_token": 1.04e-06, "input_cost_per_token": 5.2e-07, @@ -1259,6 +1369,41 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "output_cost_per_token_above_272k_tokens_priority": 4.6e-05, "cache_read_input_token_cost_above_272k_tokens": 4.6e-07, }, + "azure/gpt-5.6-terra": { + "cache_creation_input_token_cost": 2.75e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5.5e-06, + "cache_creation_input_token_cost_above_272k_tokens_flex": 2.75e-06, + "cache_creation_input_token_cost_priority": 5.5e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1.1e-05, + "cache_creation_input_token_cost_flex": 1.375e-06, + "cache_read_input_token_cost": 2.2e-07, + "cache_read_input_token_cost_above_272k_tokens": 4.4e-07, + "cache_read_input_token_cost_above_272k_tokens_flex": 2.2e-07, + "cache_read_input_token_cost_priority": 4.4e-07, + "cache_read_input_token_cost_above_272k_tokens_priority": 8.8e-07, + "cache_read_input_token_cost_flex": 1.1e-07, + "input_cost_per_token": 2.2e-06, + "input_cost_per_token_above_272k_tokens": 4.4e-06, + "input_cost_per_token_above_272k_tokens_flex": 2.2e-06, + "input_cost_per_token_priority": 4.4e-06, + "input_cost_per_token_above_272k_tokens_priority": 8.8e-06, + "input_cost_per_token_flex": 1.1e-06, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.32e-05, + "output_cost_per_token_above_272k_tokens": 1.98e-05, + "output_cost_per_token_above_272k_tokens_flex": 9.9e-06, + "output_cost_per_token_priority": 2.64e-05, + "output_cost_per_token_above_272k_tokens_priority": 3.96e-05, + "output_cost_per_token_flex": 6.6e-06, + "supports_function_calling": True, + "supports_prompt_caching": True, + "supports_reasoning": True, + "supports_none_reasoning_effort": True, + }, "azure/gpt-6-luna": { "cache_read_input_token_cost": 1.3e-08, "cache_creation_input_token_cost": 1.625e-07, From 86a91f35de4c6d6cd95510936ad1f14c288740e7 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 16 Oct 2026 17:37:33 +0000 Subject: [PATCH 07/58] fix(responses): log and bill every round of non-stream MCP auto-execute (#43667) * fix(responses): log and bill every round of non-stream MCP auto-execute Both model rounds run as internal sub-calls on the request's logging object, and the outer call logs the response the client received with usage, cost, and cost breakdown summed across the rounds. The internal-call gate now also covers the sync success callbacks, so an inner round no longer takes their first-wins slot * fix(caching): let the parent call log and bill an internal sub-call's cache hit A round of non-stream MCP auto-execute that the response cache served fired the cache-hit callbacks on the shared logging object, so the spend log kept round 1's function_call row and dropped the billed row for the answer. Internal sub-call cache hits now fire no callbacks and carry a zero cost, and the summed cost breakdown leaves out rounds served from the cache * fix(logging): keep a parent row unique when its response came from the response cache A RAG query or an MCP turn whose final round is served from the cache hands back the cached response's id, which already belongs to the row that cached it, so SpendLogs dropped the new row as a duplicate. Such a row now gets the cache-hit id suffix, and it is logged as a cache hit when it cost nothing. * fix(logging): skip the zero-cost warning for a response served from the cache at no charge A non-stream MCP auto-execute request whose rounds the response cache served is priced at $0 before the parent marks it a cache hit, so it logged a false 'priced at $0' warning and counted in litellm_zero_cost_requests_total. The diagnostic now treats a response the cache served at no charge the same as a marked cache hit * fix(responses): bill the paid first MCP round when a later step of the request fails * fix(responses): log each MCP tool call under its own id so a failed follow-up keeps its spend row The non-stream auto-execute path passed the request's call id to the tool call, so the tool's spend row and the request's failure row shared one request_id and the failure row was dropped on insert, leaving the paid first round unbilled. * fix(responses): bill round 1 on a client disconnect, keep MCP tool rows in the request's session, and charge a fixed margin once per request * fix(responses): keep bridged MCP auto-execute in one session and trace and continue its conversation The bridged follow-up no longer chains on round 1, so a thinking model keeps its reasoning and the request stays in one session. The tool call log continues the request trace without renaming it, the tools are fetched once, and a later turn finds the MCP turn under the id the client got, skipping the tool rows. * fix(responses): log the non-stream MCP tool listing in the request's session and trace * fix(responses): find a bridged MCP turn by the id the client holds and log its listing in the bridged trace * fix(responses): derive the MCP tool row id from the request's call id The tool call's spend row now uses _, so it stays linked to the request without colliding with the request's failure row. Also tighten the internal cache-hit test to a fixed clock and typed params. * fix(responses): bill a retried MCP attempt for its own rounds A router retry or fallback reuses the request's logging object, so a failed attempt's stashed round 1 usage made the next attempt's success row log stale tokens. Each non-stream MCP attempt now starts without that stash. * fix(responses): bill every failed MCP attempt's paid first round on the request failure A retry or fallback now sets a failed attempt's stashed round 1 aside instead of dropping it, and the proxy's failure row adds it to the last attempt's own, so a request whose later attempt fails on its first call still bills the round the provider charged. A cached final round also keeps the paid first round's cost breakdown. * fix(responses): keep MCP tool catalog and results out of the bridged next-turn replay * fix(responses): keep MCP tool spend, client body keys, and Langfuse traces working on auto-execute Count the MCP tools an auto-executed Responses API call ran in the tool spend rollup, accept a client body key named previous_response_client_ids without a 500, and point MCP listing and tool spans at the trace Langfuse gives the parent request, including x-litellm-session-id requests * fix(responses): join MCP spans to the Langfuse trace when the request sends body metadata * test(integration): cover non-stream MCP auto-execute billing, failure rows, and bridged follow-ups --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/_internal_context.py | 10 + litellm/caching/caching_handler.py | 11 +- litellm/integrations/langfuse/langfuse.py | 71 +- litellm/litellm_core_utils/hidden_params.py | 5 + litellm/litellm_core_utils/litellm_logging.py | 30 +- .../_experimental/mcp_server/operations.py | 2 + litellm/proxy/db/spend_log_tool_index.py | 65 +- .../spend_tracking/spend_tracking_utils.py | 4 +- litellm/proxy/utils.py | 22 + litellm/rag/main.py | 30 +- .../handler.py | 8 + .../session_handler.py | 62 +- .../transformation.py | 2 + litellm/responses/main.py | 325 ++++--- .../mcp/litellm_proxy_mcp_handler.py | 13 +- litellm/responses/mcp/round_billing.py | 134 +++ litellm/utils.py | 17 +- .../integration/mcp/test_mcp_llm_endpoints.py | 227 ++++- tests/unit/caching/test_caching_handler.py | 57 +- .../litellm_core_utils/test_hidden_params.py | 22 + .../test_litellm_logging.py | 71 ++ .../proxy/db/test_spend_log_tool_index.py | 33 + .../test_spend_tracking_utils.py | 30 + .../test_anthropic_responses_bridge.py | 1 + .../test_handler.py | 89 +- .../test_session_handler.py | 104 ++- .../mcp/test_aresponses_api_with_mcp.py | 844 +++++++++++++++++- .../unit/responses/mcp/test_round_billing.py | 218 +++++ 28 files changed, 2304 insertions(+), 203 deletions(-) create mode 100644 litellm/responses/mcp/round_billing.py create mode 100644 tests/unit/responses/mcp/test_round_billing.py diff --git a/litellm/_internal_context.py b/litellm/_internal_context.py index df549701cf4..244efc8047b 100644 --- a/litellm/_internal_context.py +++ b/litellm/_internal_context.py @@ -37,6 +37,16 @@ REDIS_FAMILIES_METADATA_KEY: Final = "families" _service_caller: Final[ContextVar[str | None]] = ContextVar("service_caller", default=None) +@contextmanager +def internal_sub_call() -> Generator[None]: + """A nested call the parent call logs and bills, so the parent must fold this call's cost into its own.""" + token: Final = is_internal_call.set(True) + try: + yield + finally: + is_internal_call.reset(token) + + @contextmanager def post_response_phase() -> Generator[None]: """Work the caller no longer waits for (success callbacks, response-cache writes), including tasks it spawns.""" diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 501dc266ba0..11488a03f3f 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -24,7 +24,7 @@ from typing import TYPE_CHECKING, Any, Final, Optional, TypeVar from pydantic import ConfigDict, SkipValidation, TypeAdapter, ValidationError import litellm -from litellm._internal_context import post_response_phase +from litellm._internal_context import is_internal_call, post_response_phase from litellm._logging import print_verbose, verbose_logger from litellm.caching import InMemoryCache from litellm.caching.caching import S3Cache, response_cache_phase @@ -511,7 +511,9 @@ class LLMCachingHandler: custom_llm_provider=custom_llm_provider, ) - if not _should_defer_streaming_cache_hit_callbacks(cached_result=cached_result): + if not is_internal_call.get() and not _should_defer_streaming_cache_hit_callbacks( + cached_result=cached_result + ): logging_obj.handle_sync_success_callbacks_for_async_calls( result=cached_result, start_time=start_time, @@ -820,7 +822,12 @@ class LLMCachingHandler: start_time (datetime): The start time of the operation. end_time (datetime): The end time of the operation. cache_hit (bool): Whether it was a cache hit. + + An internal sub-call's hit fires nothing, since the parent call that folds it in logs it. """ + if is_internal_call.get(): + return + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index 42f605443c7..86f80047e52 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -103,6 +103,43 @@ def _is_session_header_trace(trace_id: object, session_id: object, proxy_server_ ) +def langfuse_trace_metadata(metadata: object, proxy_server_request: object) -> Mapping[str, object]: + entries: Final = _object_mapping(metadata) or _NO_METADATA + request: Final = _object_mapping(proxy_server_request) + headers: Final = _object_mapping(request.get("headers")) if request is not None else None + if headers is None: + return entries + header_entries: Final = { + key.removeprefix("langfuse_"): value + for key, value in _widened_items(headers) + if isinstance(key, str) and key.startswith("langfuse_") + } + return MappingProxyType({**entries, **header_entries}) + + +def langfuse_trace_id( + metadata: Mapping[str, object], + standard_trace_id: object, + litellm_trace_id: object, + litellm_call_id: object, + proxy_server_request: object, +) -> object: + existing_trace_id: Final = metadata.get("existing_trace_id") + if existing_trace_id is not None: + return existing_trace_id + metadata_trace_id: Final = metadata.get("trace_id") + trace_id: Final = ( + metadata_trace_id + if metadata_trace_id is not None + else standard_trace_id + if standard_trace_id is not None + else litellm_trace_id or litellm_call_id + ) + if _is_session_header_trace(trace_id, metadata.get("session_id"), proxy_server_request): + return litellm_call_id or trace_id + return trace_id + + class _UsageObject(Protocol): """Token-count surface the Langfuse logger reads off a response usage payload.""" @@ -680,33 +717,19 @@ class LangFuseLogger: # Add default langfuse tags tags = self.add_default_langfuse_tags(tags=tags, kwargs=kwargs, metadata=metadata) + call_trace_id: Final = langfuse_trace_id( + metadata=clean_metadata, + standard_trace_id=standard_logging_object.get("trace_id") + if standard_logging_object is not None + else None, + litellm_trace_id=kwargs.get("litellm_trace_id"), + litellm_call_id=litellm_call_id, + proxy_server_request=litellm_params.get("proxy_server_request"), + ) session_id: Final = clean_metadata.pop("session_id", None) trace_name = cast(str | None, clean_metadata.pop("trace_name", None)) - trace_id = clean_metadata.pop("trace_id", None) - # Use standard_logging_object.trace_id if available (when trace_id from metadata is None) - # This allows standard trace_id to be used when provided in standard_logging_object - if trace_id is None and standard_logging_object is not None: - trace_id = cast(str | None, standard_logging_object.get("trace_id")) - # Fallback to litellm_call_id if no trace_id found - if trace_id is None: - trace_id = kwargs.get("litellm_trace_id") or litellm_call_id + clean_metadata.pop("trace_id", None) existing_trace_id: Final = clean_metadata.pop("existing_trace_id", None) - # If existing_trace_id is provided, use it as the trace_id to return - # This allows continuing an existing trace while still returning the correct trace_id - if existing_trace_id is not None: - trace_id = existing_trace_id - call_trace_id: Final = ( - litellm_call_id or trace_id - if existing_trace_id is None - and _is_session_header_trace(trace_id, session_id, litellm_params.get("proxy_server_request")) - else trace_id - ) - if call_trace_id != trace_id: - verbose_logger.debug( - "Langfuse: trace_id %s came from a session header; using call id %s so each call gets its own trace", - trace_id, - call_trace_id, - ) requested_trace_keys: Final = _as_steering_key_sequence(clean_metadata.pop("update_trace_keys", ())) update_trace_keys: Final = ( requested_trace_keys if _as_steering_flag(litellm.langfuse_enable_update_trace_keys) else () diff --git a/litellm/litellm_core_utils/hidden_params.py b/litellm/litellm_core_utils/hidden_params.py index 714b258413e..4cb9dcd53f1 100644 --- a/litellm/litellm_core_utils/hidden_params.py +++ b/litellm/litellm_core_utils/hidden_params.py @@ -77,6 +77,11 @@ def get_hidden_params(obj: object) -> MutableMapping[str, object] | None: return _as_hidden_params_mapping(hidden_params) +def served_from_cache(obj: object) -> bool: + hidden_params: Final = get_hidden_params(obj) + return hidden_params is not None and hidden_params.get("cache_hit") is True + + def set_hidden_params(obj: object, hidden_params: dict[str, object] | HiddenParams) -> None: if isinstance(obj, dict): obj[_HIDDEN_PARAMS_ATTR] = hidden_params diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index fa009ca6327..67d1573dbc4 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -79,7 +79,12 @@ from litellm.litellm_core_utils.core_helpers import ( ) from litellm.litellm_core_utils.error_normalization import normalize_error from litellm.litellm_core_utils.get_litellm_params import get_litellm_params -from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_param +from litellm.litellm_core_utils.hidden_params import ( + HIDDEN_PARAMS_ATTR, + get_hidden_params, + served_from_cache, + set_hidden_param, +) from litellm.litellm_core_utils.internal_call_metadata import ( MODEL_ACCESS_GROUP_METADATA_KEY, is_unbilled_non_inference_call, @@ -2056,7 +2061,7 @@ class Logging(LiteLLMLoggingBaseClass): ) -> None: if response_cost is None and not calculation_failed: return - if self.model_call_details.get("cache_hit") is True: + if self.model_call_details.get("cache_hit") is True or _served_from_cache_at_no_charge(result): self.model_call_details["zero_cost_diagnostic"] = None return try: @@ -2295,6 +2300,17 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["combined_usage_object"] = usage self.model_call_details["response_cost"] = response_cost + def carry_failed_attempt_usage(self) -> None: + """Move what a failed earlier attempt stashed off the retry sharing this logging object, so the retry logs its + own usage while the request's failure log still bills every attempt.""" + usage: Final = self.model_call_details.pop("combined_usage_object", None) + response_cost: Final = self.model_call_details.pop("response_cost", None) + if isinstance(usage, Usage): + self.model_call_details["failed_attempts_usage"] = ( + *self.model_call_details.get("failed_attempts_usage", ()), + (usage, float(response_cost) if isinstance(response_cost, (int, float)) else 0.0), + ) + def record_continued_partial(self, usage: Usage, response_cost: float, response_id: str) -> None: """Bill what a cut stream already consumed on this continuation's success log, under the id the caller saw.""" self.model_call_details["continued_partial_usage"] = usage @@ -2662,6 +2678,8 @@ class Logging(LiteLLMLoggingBaseClass): result = self._handle_a2a_response_logging(result=result) logging_result: Final = self.normalize_logging_result(result=result) + if cache_hit is None and _served_from_cache_at_no_charge(logging_result): + self.model_call_details["cache_hit"] = True if isinstance(result, Response) and isinstance( logging_result, (ModelResponse, EmbeddingResponse, TextCompletionResponse) @@ -6679,6 +6697,11 @@ def _get_status_fields( return StandardLoggingPayloadStatusFields(llm_api_status=llm_api_status, guardrail_status=guardrail_status) +def _served_from_cache_at_no_charge(result: object) -> bool: + hidden_params: Final = get_hidden_params(result) + return served_from_cache(result) and hidden_params is not None and hidden_params.get("response_cost") == 0 + + def _extract_response_obj_and_hidden_params( init_response_obj: object, original_exception: Exception | None, @@ -6841,8 +6864,9 @@ def get_standard_logging_object_payload( ) # maintain backwards compatibility with old request body check saved_cache_cost: float = 0.0 - if cache_hit is True: + if cache_hit is True or served_from_cache(init_response_obj): id = f"{id}_cache_hit{time.time()}" # do not duplicate the request id + if cache_hit is True: saved_cache_cost = ( logging_obj.response_cost_calculator( result=init_response_obj, diff --git a/litellm/proxy/_experimental/mcp_server/operations.py b/litellm/proxy/_experimental/mcp_server/operations.py index 2de54e8ce1b..90f01368127 100644 --- a/litellm/proxy/_experimental/mcp_server/operations.py +++ b/litellm/proxy/_experimental/mcp_server/operations.py @@ -1160,6 +1160,7 @@ async def _get_tools_from_mcp_servers( protocol_version: str | None = None, *, record_listing: bool = False, + request_trace_id: str | None = None, enforce_rate_limits: bool = True, ) -> AggregateToolListing: """ @@ -1206,6 +1207,7 @@ async def _get_tools_from_mcp_servers( "spend_logs_metadata": spend_logs_metadata, "headers": logging_safe_mcp_headers(raw_headers), **({"tags": request_tags} if request_tags else {}), + **({"existing_trace_id": request_trace_id} if request_trace_id else {}), }, # Provide a small input payload for standard logging "input": [ diff --git a/litellm/proxy/db/spend_log_tool_index.py b/litellm/proxy/db/spend_log_tool_index.py index 81338988c66..7ac414bec00 100644 --- a/litellm/proxy/db/spend_log_tool_index.py +++ b/litellm/proxy/db/spend_log_tool_index.py @@ -18,9 +18,10 @@ import random from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timezone -from itertools import groupby -from typing import TYPE_CHECKING, Any, Final, Protocol +from itertools import chain, groupby +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from typing_extensions import LiteralString from litellm.constants import SPEND_LOG_WRITE_BATCH_MAX_BYTES, SPEND_LOG_WRITE_BATCH_MAX_ROWS @@ -47,21 +48,73 @@ class ToolUsageTransaction: def response_tool_call_names(completion_response: object) -> tuple[str, ...]: """Tool names invoked in a completion response, in call order, for any response surface get_tool_calls_from_response understands (chat completions, Responses - API output items, Anthropic Messages tool_use blocks). Reads every choice of - an ``n>1`` chat response: each choice cost money and its tool calls ran.""" + API output items, Anthropic Messages tool_use blocks), plus the MCP tools an + auto-executed Responses API call ran. Reads every choice of an ``n>1`` chat + response: each choice cost money and its tool calls ran.""" if completion_response is None or isinstance(completion_response, Exception): return () from litellm.litellm_core_utils.prompt_templates.factory import ( get_tool_calls_from_response, ) + called_names: Final = ( + tool_call.get("name") + for tool_call in get_tool_calls_from_response(completion_response, include_all_choices=True) + ) return tuple( stripped - for tool_call in get_tool_calls_from_response(completion_response, include_all_choices=True) - if isinstance(name := tool_call.get("name"), str) and (stripped := name.strip()) + for name in chain(called_names, _auto_executed_mcp_tool_names(completion_response)) + if isinstance(name, str) and (stripped := name.strip()) ) +class _ExecutedTool(BaseModel): + model_config = ConfigDict(frozen=True) + + name: str | None = None + + +class _OutputText(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + text: str + + +class _ToolExecutionResults(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + type: Literal["tool_execution_results"] + content: tuple[_OutputText, ...] + + +class _ResponsesApiOutput(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + output: tuple[object, ...] + + +_EXECUTED_TOOLS: Final = TypeAdapter(tuple[_ExecutedTool, ...]) + + +def _auto_executed_mcp_tool_names(completion_response: object) -> tuple[str | None, ...]: + try: + response: Final = _ResponsesApiOutput.model_validate(completion_response) + except ValidationError: + return () + return tuple(chain.from_iterable(_executed_tool_names(item) for item in response.output)) + + +def _executed_tool_names(output_item: object) -> tuple[str | None, ...]: + try: + results: Final = _ToolExecutionResults.model_validate(output_item) + executed: Final = tuple( + chain.from_iterable(_EXECUTED_TOOLS.validate_json(content.text) for content in results.content) + ) + except ValidationError: + return () + return tuple(tool.name for tool in executed) + + def build_tool_usage_transaction( request_id: str, start_time_iso: str, diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 9ebea589cc9..c7ca08fb513 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -47,6 +47,7 @@ from litellm.litellm_core_utils.core_helpers import ( reconstruct_model_name, ) from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider +from litellm.litellm_core_utils.hidden_params import served_from_cache from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call from litellm.litellm_core_utils.litellm_logging import ( coerce_model_access_groups, @@ -542,6 +543,7 @@ def get_logging_payload( kwargs = {} rejected_as_unknown_model: Final = isinstance(response_obj, ProxyModelNotFoundError) + response_served_from_cache: Final = served_from_cache(response_obj) if response_obj is None: response_obj = {} elif not isinstance(response_obj, BaseModel) and not isinstance(response_obj, dict): @@ -812,7 +814,7 @@ def get_logging_payload( cache_key = litellm_params["preset_cache_key"] else: cache_key = litellm.cache.get_cache_key(**kwargs) - if cache_hit is True: + if cache_hit is True or response_served_from_cache: import time id = f"{id}_cache_hit{time.time()}" # SpendLogs does not allow duplicate request_id diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index be6532ba36b..dc6423c979b 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -65,6 +65,7 @@ from litellm.constants import ( SPEND_LOG_WRITE_BATCH_MAX_BYTES, SPEND_LOG_WRITE_BATCH_MAX_ROWS, ) +from litellm.cost_calculator import BaseTokenUsageProcessor from litellm.litellm_core_utils.bug_report import ( bug_report_notice, should_report_bug, @@ -990,6 +991,24 @@ _INPUT_ESTIMABLE_CALL_TYPES: Final = frozenset( ) +_FAILED_ATTEMPTS_USAGE: Final = TypeAdapter(tuple[tuple[Usage, float], ...]) + + +def _billed_across_attempts( + failed_attempts: tuple[tuple[Usage, float], ...], recovered_usage: object, recovered_cost: object +) -> tuple[Usage, float]: + this_attempt: Final = ( + ((recovered_usage, float(recovered_cost) if isinstance(recovered_cost, (int, float)) else 0.0),) + if isinstance(recovered_usage, Usage) + else () + ) + attempts: Final = (*failed_attempts, *this_attempt) + return ( + BaseTokenUsageProcessor.combine_usage_objects([usage for usage, _ in attempts]), + sum(response_cost for _, response_cost in attempts), + ) + + def _failure_usage_to_lift( model_call_details: Mapping[str, object], request_body: Mapping[str, object], @@ -1004,6 +1023,9 @@ def _failure_usage_to_lift( chat-completions provider) without filling optional_params. Returns the (combined_usage_object, response_cost) pair to lift, or None.""" recovered_usage: Final = model_call_details.get("combined_usage_object") + failed_attempts: Final = _FAILED_ATTEMPTS_USAGE.validate_python(model_call_details.get("failed_attempts_usage", ())) + if failed_attempts: + return _billed_across_attempts(failed_attempts, recovered_usage, model_call_details.get("response_cost")) if recovered_usage is not None: return recovered_usage, model_call_details.get("response_cost") if not dispatched or model_call_details.get(LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL): diff --git a/litellm/rag/main.py b/litellm/rag/main.py index 5ab39bec080..be566a0da8b 100644 --- a/litellm/rag/main.py +++ b/litellm/rag/main.py @@ -11,8 +11,7 @@ __all__ = ["aingest", "aquery", "ingest", "query"] import asyncio import contextvars -from collections.abc import Coroutine, Iterator, Mapping -from contextlib import contextmanager +from collections.abc import Coroutine, Mapping from functools import partial from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final @@ -21,7 +20,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter import litellm -from litellm._internal_context import is_internal_call +from litellm._internal_context import internal_sub_call from litellm.cost_calculator import vector_store_search_cost from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion @@ -212,25 +211,6 @@ async def aingest( ) -@contextmanager -def _suppressed_sub_call_billing() -> Iterator[None]: - """ - Suppress a sub-call's own billing event so the parent aquery event bills it. - - Every suppressed sub-call's cost must be folded into the parent event: - into the response's hidden response_cost on the non-streaming path, or via - the logging object's additional_response_cost on the streaming path (the - streamed cost is computed from assembled chunks after this pipeline - returns, so there is no response object to fold into here). - """ - previous: Final = is_internal_call.get() - is_internal_call.set(True) - try: - yield - finally: - is_internal_call.set(previous) - - async def _execute_query_pipeline( model: str, messages: list[AllMessageValues], @@ -272,7 +252,7 @@ async def _execute_query_pipeline( forwarded_search_params: Final = MappingProxyType( {**provider_search_params, **kwargs, **filter_search_params, **store_search_params} ) - with _suppressed_sub_call_billing(): + with internal_sub_call(): search_response: Final = await litellm.vector_stores.asearch( vector_store_id=retrieval_config["vector_store_id"], query=query_text, @@ -306,7 +286,7 @@ async def _execute_query_pipeline( if rerank and rerank.get("enabled"): documents: Final = RAGQuery.extract_documents_from_search(search_response) if documents: - with _suppressed_sub_call_billing(): + with internal_sub_call(): rerank_response = await litellm.arerank( model=rerank["model"], query=query_text, @@ -324,7 +304,7 @@ async def _execute_query_pipeline( modified_messages: Final = messages[:-1] + [context_message] + [messages[-1]] # Use router if available to properly resolve virtual model names - with _suppressed_sub_call_billing(): + with internal_sub_call(): if router is not None: response = await router.acompletion( model=model, diff --git a/litellm/responses/litellm_completion_transformation/handler.py b/litellm/responses/litellm_completion_transformation/handler.py index 0ab2052abb0..5489d2223ac 100644 --- a/litellm/responses/litellm_completion_transformation/handler.py +++ b/litellm/responses/litellm_completion_transformation/handler.py @@ -29,6 +29,9 @@ _REQUEST_KWARGS_ADAPTER: Final = TypeAdapter(dict[str, object]) class LiteLLMCompletionTransformationHandler: def response_api_handler( self, + previous_response_client_ids: tuple[str, ...] = (), + /, + *, model: str, input: str | ResponseInputParam, responses_api_request: ResponsesAPIOptionalRequestParams, @@ -61,6 +64,7 @@ class LiteLLMCompletionTransformationHandler: if _is_async: return self.async_response_api_handler( + previous_response_client_ids, litellm_completion_request=litellm_completion_request, request_input=bridged_input, responses_api_request=bridged_request, @@ -100,6 +104,9 @@ class LiteLLMCompletionTransformationHandler: async def async_response_api_handler( self, + previous_response_client_ids: tuple[str, ...] = (), + /, + *, litellm_completion_request: dict, request_input: str | ResponseInputParam, responses_api_request: ResponsesAPIOptionalRequestParams, @@ -112,6 +119,7 @@ class LiteLLMCompletionTransformationHandler: litellm_completion_request=litellm_completion_request, instructions=responses_api_request.get("instructions"), request_kwargs=_REQUEST_KWARGS_ADAPTER.validate_python(kwargs), + previous_response_client_ids=previous_response_client_ids, ) acompletion_args: Final = {} diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index e78ec5e5d3e..615bdb5bbfb 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -14,7 +14,7 @@ from litellm.types.llms.openai import ( GenericChatCompletionMessage, ResponseInputParam, ) -from litellm.types.utils import ChatCompletionMessageToolCall, Message, ModelResponse +from litellm.types.utils import CallTypes, ChatCompletionMessageToolCall, Message, ModelResponse if TYPE_CHECKING: from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload @@ -42,6 +42,30 @@ def _normalize_redacted_tool_call_arguments(message: Message) -> None: function_call.arguments = REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER +_MCP_TOOL_CALL_TYPES: Final = frozenset((CallTypes.list_mcp_tools.value, CallTypes.call_mcp_tool.value)) + + +def _is_a_conversation_turn(spend_log: "SpendLogsPayload") -> bool: + return spend_log.get("call_type") not in _MCP_TOOL_CALL_TYPES + + +_MCP_AUTO_EXECUTE_OUTPUT_TYPES: Final = frozenset(("mcp_tools_fetched", "tool_execution_results")) + + +def _is_model_output(item: object) -> bool: + return not (isinstance(item, Mapping) and item.get("type") in _MCP_AUTO_EXECUTE_OUTPUT_TYPES) + + +def _stored_responses_api_output(stored_response: object) -> ResponseInputParam | None: + match stored_response: + case {"object": "response", "output": list() as output}: + return [item for item in output if _is_model_output(item)] + case {"object": "response"}: + return [] + case _: + return None + + def _stored_instructions(proxy_server_request: Mapping[str, object] | None) -> str | None: instructions: Final = None if proxy_server_request is None else proxy_server_request.get("instructions") return instructions if isinstance(instructions, str) and instructions else None @@ -52,6 +76,7 @@ class ResponsesSessionHandler: async def get_chat_completion_message_history_for_previous_response_id( previous_response_id: str, item_cache_control: bool = False, + client_response_ids: tuple[str, ...] = (), ) -> ChatCompletionSession: """ Return the chat completion message history for a previous response id @@ -61,9 +86,13 @@ class ResponsesSessionHandler: ) verbose_proxy_logger.debug("inside get_chat_completion_message_history_for_previous_response_id") - all_spend_logs: list[ - SpendLogsPayload - ] = await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id(previous_response_id) + all_spend_logs: Final = [ + spend_log + for spend_log in await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id( + previous_response_id, client_response_ids + ) + if _is_a_conversation_turn(spend_log) + ] verbose_proxy_logger.debug("found %s spend logs for this response id", len(all_spend_logs)) litellm_session_id: str | None = None @@ -158,8 +187,17 @@ class ResponsesSessionHandler: ############################################################ # Add Output messages for this Spend Log ############################################################ - _response_output: Final = spend_log.get("response", "{}") - if isinstance(_response_output, dict) and _response_output and _response_output != {}: + _response_output: Final[object] = spend_log.get("response", "{}") + if (responses_api_output := _stored_responses_api_output(_response_output)) is not None: + chat_completion_message_history.extend( + LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=responses_api_output, + responses_api_request={}, + replay_reasoning=True, + item_cache_control=item_cache_control, + ) + ) + elif isinstance(_response_output, dict) and _response_output and _response_output != {}: # transform `ChatCompletion Response` to `ResponsesAPIResponse` model_response: Final = ModelResponse(**_response_output) for choice in model_response.choices: @@ -176,7 +214,7 @@ class ResponsesSessionHandler: Get the parsed proxy server request from the spend log """ proxy_server_request: Final[str | dict] = spend_log.get("proxy_server_request") or "{}" - proxy_server_request_dict: dict | None = None + proxy_server_request_dict: dict[str, object] | None = None if isinstance(proxy_server_request, dict): proxy_server_request_dict = proxy_server_request else: @@ -229,7 +267,7 @@ class ResponsesSessionHandler: @staticmethod async def get_proxy_server_request_from_cold_storage_with_object_key( object_key: str, - ) -> dict | None: + ) -> dict[str, object] | None: """ Get the proxy server request from cold storage using the object key directly. @@ -251,7 +289,7 @@ class ResponsesSessionHandler: @staticmethod def _should_check_cold_storage_for_full_payload( - proxy_server_request_dict: dict | None, + proxy_server_request_dict: Mapping[str, object] | None, ) -> bool: """ Only check cold storage when both are true @@ -274,6 +312,7 @@ class ResponsesSessionHandler: @staticmethod async def get_all_spend_logs_for_previous_response_id( previous_response_id: str, + client_response_ids: tuple[str, ...] = (), ) -> "list[SpendLogsPayload]": """ Get all spend logs for a previous response id @@ -307,11 +346,12 @@ class ResponsesSessionHandler: if prisma_client is None: return [] + request_ids: Final = list(dict.fromkeys((response_id, previous_response_id, *client_response_ids))) query: Final = """ WITH matching_session AS ( SELECT session_id, api_key, "user", team_id FROM "LiteLLM_SpendLogs" - WHERE request_id = $1 + WHERE request_id = ANY($1::text[]) ) SELECT * FROM "LiteLLM_SpendLogs" AS logs @@ -334,7 +374,7 @@ class ResponsesSessionHandler: for attempt in range(max_attempts): if attempt: await asyncio.sleep(RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL) - if spend_logs := await prisma_client.db.query_raw(query, response_id): + if spend_logs := await prisma_client.db.query_raw(query, request_ids): verbose_proxy_logger.debug( "Found the following spend logs for previous response id %s: %s", response_id, diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 587cdfcee25..35d4b8bf1ba 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -1008,6 +1008,7 @@ class LiteLLMCompletionResponsesConfig: instructions: str | None, session_lookup: Callable[[str], Awaitable[ChatCompletionSession]] | None = None, request_kwargs: Mapping[str, object] = MappingProxyType({}), + previous_response_client_ids: tuple[str, ...] = (), ) -> dict: """ Async hook to get the chain of previous input and output pairs and return a list of Chat Completion messages @@ -1022,6 +1023,7 @@ class LiteLLMCompletionResponsesConfig: lookup: Final = session_lookup or partial( ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id, item_cache_control=keeps_item_marks, + client_response_ids=previous_response_client_ids, ) chat_completion_session: Final = ( await lookup(previous_response_id) diff --git a/litellm/responses/main.py b/litellm/responses/main.py index f566c8f4e02..b7385c3b3d5 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -16,6 +16,7 @@ from typing_extensions import ( ) import litellm +from litellm._internal_context import internal_sub_call from litellm._logging import verbose_logger from litellm.completion_extras.litellm_responses_transformation.transformation import ( LiteLLMResponsesTransformationHandler, @@ -29,6 +30,7 @@ from litellm.integrations.anthropic_cache_control_hook import ( from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.core_helpers import normalize_drop_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.litellm_core_utils.prompt_templates.common_utils import ( update_responses_input_with_model_file_ids, update_responses_tools_with_model_file_ids, @@ -41,6 +43,11 @@ from litellm.responses.litellm_completion_transformation.handler import ( LiteLLMCompletionTransformationHandler, ) from litellm.responses.mcp.request_context import MCPRequestContext +from litellm.responses.mcp.round_billing import ( + billed_for_every_round, + first_round_billed_on_failure, + summed_cost_breakdown, +) from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ( PromptObject, @@ -64,7 +71,7 @@ from litellm.llms.openai.data_residency import infer_openai_data_residency from litellm.types.responses.main import * from litellm.types.responses.streaming_websocket import ResponsesWebSocketRequestDefaults from litellm.types.router import GenericLiteLLMParams -from litellm.types.utils import all_litellm_params +from litellm.types.utils import ADDRESSED_RESPONSE_ID_FIELD, all_litellm_params from litellm.utils import ( ProviderConfigManager, client, @@ -161,6 +168,67 @@ def mock_responses_api_response( ) +_STR_MAPPING: Final = TypeAdapter(Mapping[str, object]) + + +def _client_trace_id(request_kwargs: Mapping[str, object]) -> str | None: + trace_id: Final = request_kwargs.get("litellm_trace_id") + return trace_id if isinstance(trace_id, str) else None + + +def _request_call_id(request_kwargs: Mapping[str, object]) -> str | None: + call_id: Final = request_kwargs.get("litellm_call_id") + return call_id if isinstance(call_id, str) else None + + +def _logging_params(logging_obj: LiteLLMLoggingObj) -> Mapping[str, object]: + return _STR_MAPPING.validate_python(logging_obj.litellm_params) + + +def _request_logging_params( + logging_obj: LiteLLMLoggingObj, request_kwargs: Mapping[str, object] +) -> Mapping[str, object]: + logged_metadata: Final = ( + request_kwargs["litellm_metadata"] if "litellm_metadata" in request_kwargs else request_kwargs.get("metadata") + ) + return {**_logging_params(logging_obj), **request_kwargs, "metadata": logged_metadata} + + +def _request_trace_id(request_kwargs: Mapping[str, object]) -> str | None: + logging_obj: Final = request_kwargs.get("litellm_logging_obj") + if isinstance(logging_obj, LiteLLMLoggingObj): + return StandardLoggingPayloadSetup.get_standard_logging_payload_trace_id( + logging_obj=logging_obj, litellm_params=_logging_params(logging_obj) + ) + return _client_trace_id(request_kwargs) + + +def _langfuse_request_trace_id(request_kwargs: Mapping[str, object], standard_trace_id: str | None) -> str | None: + from litellm.integrations.langfuse.langfuse import langfuse_trace_id, langfuse_trace_metadata + + logging_obj: Final = request_kwargs.get("litellm_logging_obj") + if not isinstance(logging_obj, LiteLLMLoggingObj): + return None + litellm_params: Final = _request_logging_params(logging_obj, request_kwargs) + trace_id: Final = langfuse_trace_id( + metadata=langfuse_trace_metadata(litellm_params.get("metadata"), litellm_params.get("proxy_server_request")), + standard_trace_id=standard_trace_id, + litellm_trace_id=logging_obj.litellm_trace_id, + litellm_call_id=logging_obj.litellm_call_id, + proxy_server_request=litellm_params.get("proxy_server_request"), + ) + return trace_id if isinstance(trace_id, str) else None + + +def _listing_trace_id(request_kwargs: Mapping[str, object], *, bridged_first_turn: bool) -> str | None: + logging_obj: Final = request_kwargs.get("litellm_logging_obj") + if bridged_first_turn and isinstance(logging_obj, LiteLLMLoggingObj): + return StandardLoggingPayloadSetup.get_standard_logging_payload_trace_id( + logging_obj=logging_obj, litellm_params=_request_logging_params(logging_obj, request_kwargs) + ) + return _request_trace_id(request_kwargs) + + async def aresponses_api_with_mcp( input: str | ResponseInputParam, model: str, @@ -223,6 +291,22 @@ async def aresponses_api_with_mcp( else (None, None, None, None) ) + request_kwargs: Final = _STR_MAPPING.validate_python(kwargs) + logging_obj: Final = request_kwargs.get("litellm_logging_obj") + joins_request_trace: Final = isinstance(logging_obj, LiteLLMLoggingObj) + bridged_first_turn: Final = ( + joins_request_trace + and not stream + and previous_response_id is None + and not _provider_stores_the_first_round(model, custom_llm_provider, kwargs) + ) + + listing_trace_id: Final = ( + _client_trace_id(request_kwargs) + if stream + else _listing_trace_id(request_kwargs, bridged_first_turn=bridged_first_turn) + ) + # Get original MCP tools (for events) and OpenAI tools (for LLM) by reusing existing methods ( original_mcp_tools, @@ -230,7 +314,8 @@ async def aresponses_api_with_mcp( ) = await LiteLLM_Proxy_MCP_Handler.process_mcp_tools_without_openai_transform( user_api_key_auth=user_api_key_auth, mcp_tools_with_litellm_proxy=mcp_tools_with_litellm_proxy, - litellm_trace_id=kwargs.get("litellm_trace_id"), + litellm_trace_id=listing_trace_id, + request_trace_id=None if stream else _langfuse_request_trace_id(request_kwargs, listing_trace_id), mcp_auth_header=mcp_auth_header, mcp_server_auth_headers=mcp_server_auth_headers, request_tags=LiteLLM_Proxy_MCP_Handler.get_parent_request_tags(kwargs), @@ -312,13 +397,17 @@ async def aresponses_api_with_mcp( ######################################################### # Make initial response API call ######################################################### - response: Final = await aresponses( - input=input, - model=model, - tools=all_tools, - previous_response_id=previous_response_id, - **initial_call_params, - ) + if isinstance(logging_obj, LiteLLMLoggingObj): + logging_obj.carry_failed_attempt_usage() + with internal_sub_call(): + response: Final = await aresponses( + input=input, + model=model, + tools=all_tools, + previous_response_id=previous_response_id, + **initial_call_params, + ) + first_round_breakdown: Final = logging_obj.cost_breakdown if isinstance(logging_obj, LiteLLMLoggingObj) else None verbose_logger.debug("Initial response %s", response) @@ -327,107 +416,116 @@ async def aresponses_api_with_mcp( # If auto-execute tools is True, then we need to execute the tool calls ######################################################### if should_auto_execute and isinstance(response, ResponsesAPIResponse): - tool_calls: Final = LiteLLM_Proxy_MCP_Handler.extract_tool_calls_from_response(response=response) + with first_round_billed_on_failure(logging_obj, response): + tool_calls: Final = LiteLLM_Proxy_MCP_Handler.extract_tool_calls_from_response(response=response) - if tool_calls: - user_api_key_auth = kwargs.get("litellm_metadata", {}).get("user_api_key_auth") + if tool_calls: + user_api_key_auth = kwargs.get("litellm_metadata", {}).get("user_api_key_auth") - # Extract MCP auth headers from the request to pass to MCP server - ( - mcp_auth_header, - mcp_server_auth_headers, - oauth2_headers, - raw_headers_from_request, - ) = ResponsesAPIRequestUtils.extract_mcp_headers_from_request( - secret_fields=secret_fields, - tools=tools, - ) - - tool_results: Final = await LiteLLM_Proxy_MCP_Handler.execute_tool_calls( - tool_server_map=tool_server_map, - served_tools=original_mcp_tools, - tool_calls=tool_calls, - user_api_key_auth=user_api_key_auth, - mcp_auth_header=mcp_auth_header, - mcp_server_auth_headers=mcp_server_auth_headers, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers_from_request, - litellm_call_id=kwargs.get("litellm_call_id"), - litellm_trace_id=kwargs.get("litellm_trace_id"), - request_tags=LiteLLM_Proxy_MCP_Handler.get_parent_request_tags(kwargs), - guardrail_context=MCPRequestContext.resolve_guardrail_context( - MappingProxyType({**kwargs, "metadata": metadata, "model": model}) - ), - ) - - if tool_results: - persistence_disabled: Final = LiteLLM_Proxy_MCP_Handler.is_persistence_disabled(call_params) - - follow_up_input: Final = LiteLLM_Proxy_MCP_Handler.create_follow_up_input( - response=response, - tool_results=tool_results, - original_input=input, - preserve_reasoning=persistence_disabled, + # Extract MCP auth headers from the request to pass to MCP server + ( + mcp_auth_header, + mcp_server_auth_headers, + oauth2_headers, + raw_headers_from_request, + ) = ResponsesAPIRequestUtils.extract_mcp_headers_from_request( + secret_fields=secret_fields, + tools=tools, ) - # Prepare parameters for follow-up call (restores original stream setting) - follow_up_call_params: Final = LiteLLM_Proxy_MCP_Handler.prepare_follow_up_call_params( - call_params=call_params, original_stream_setting=stream or False + tool_trace_id: Final = _request_trace_id(request_kwargs) + tool_results: Final = await LiteLLM_Proxy_MCP_Handler.execute_tool_calls( + tool_server_map=tool_server_map, + served_tools=original_mcp_tools, + tool_calls=tool_calls, + user_api_key_auth=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers_from_request, + litellm_call_id=_request_call_id(request_kwargs), + litellm_trace_id=tool_trace_id, + request_tags=LiteLLM_Proxy_MCP_Handler.get_parent_request_tags(kwargs), + guardrail_context=MCPRequestContext.resolve_guardrail_context( + MappingProxyType({**kwargs, "metadata": metadata, "model": model}) + ), + request_trace_id=_langfuse_request_trace_id(request_kwargs, tool_trace_id), ) - # Create tool execution events for streaming if needed - tool_execution_events = [] - if stream: - tool_execution_events = LiteLLM_Proxy_MCP_Handler.create_tool_execution_events( - tool_calls=tool_calls, tool_results=tool_results - ) + if tool_results: + chains_on_first_round: Final = not LiteLLM_Proxy_MCP_Handler.is_persistence_disabled( + call_params + ) and _provider_stores_the_first_round(model, custom_llm_provider, kwargs) - final_response = await LiteLLM_Proxy_MCP_Handler.make_follow_up_call( - follow_up_input=follow_up_input, - model=model, - all_tools=all_tools, - response_id=previous_response_id if persistence_disabled else response.id, - **follow_up_call_params, - ) - - # If streaming and we have tool execution events, wrap the response - if ( - stream - and tool_execution_events - and (hasattr(final_response, "__aiter__") or hasattr(final_response, "__iter__")) - ): - from litellm.responses.mcp.mcp_streaming_iterator import ( - MCPEnhancedStreamingIterator, - ) - - final_response = MCPEnhancedStreamingIterator( - tool_server_map=tool_server_map, - served_tools=original_mcp_tools, - base_iterator=final_response, - mcp_events=tool_execution_events, - user_api_key_auth=user_api_key_auth, - ) - - # Add custom output elements to the final response (for non-streaming) - elif isinstance(final_response, ResponsesAPIResponse): - # Fetch MCP tools again for output elements (without OpenAI transformation) - ( - mcp_tools_for_output, - _, - ) = await LiteLLM_Proxy_MCP_Handler.process_mcp_tools_without_openai_transform( - user_api_key_auth=user_api_key_auth, - mcp_tools_with_litellm_proxy=mcp_tools_with_litellm_proxy, - mcp_auth_header=mcp_auth_header, - mcp_server_auth_headers=mcp_server_auth_headers, - request_tags=LiteLLM_Proxy_MCP_Handler.get_parent_request_tags(kwargs), - raw_headers=discovery_raw_headers, - ) - final_response = LiteLLM_Proxy_MCP_Handler.add_mcp_output_elements_to_response( - response=final_response, - mcp_tools_fetched=mcp_tools_for_output, + follow_up_input: Final = LiteLLM_Proxy_MCP_Handler.create_follow_up_input( + response=response, tool_results=tool_results, + original_input=input, + preserve_reasoning=not chains_on_first_round, ) - return final_response + + # Prepare parameters for follow-up call (restores original stream setting) + follow_up_call_params: Final = LiteLLM_Proxy_MCP_Handler.prepare_follow_up_call_params( + call_params=call_params, original_stream_setting=stream or False + ) + + # Create tool execution events for streaming if needed + tool_execution_events = [] + if stream: + tool_execution_events = LiteLLM_Proxy_MCP_Handler.create_tool_execution_events( + tool_calls=tool_calls, tool_results=tool_results + ) + + with internal_sub_call(): + final_response = await LiteLLM_Proxy_MCP_Handler.make_follow_up_call( + follow_up_input=follow_up_input, + model=model, + all_tools=all_tools, + response_id=response.id if chains_on_first_round else previous_response_id, + **follow_up_call_params, + ) + + # If streaming and we have tool execution events, wrap the response + if ( + stream + and tool_execution_events + and (hasattr(final_response, "__aiter__") or hasattr(final_response, "__iter__")) + ): + from litellm.responses.mcp.mcp_streaming_iterator import ( + MCPEnhancedStreamingIterator, + ) + + final_response = MCPEnhancedStreamingIterator( + tool_server_map=tool_server_map, + served_tools=original_mcp_tools, + base_iterator=final_response, + mcp_events=tool_execution_events, + user_api_key_auth=user_api_key_auth, + ) + + # Add custom output elements to the final response (for non-streaming) + elif isinstance(final_response, ResponsesAPIResponse): + final_round_breakdown: Final = ( + logging_obj.cost_breakdown if isinstance(logging_obj, LiteLLMLoggingObj) else None + ) + if isinstance(logging_obj, LiteLLMLoggingObj): + logging_obj.cost_breakdown = summed_cost_breakdown( + first=response, + first_breakdown=first_round_breakdown, + final=final_response, + final_breakdown=final_round_breakdown, + ) + return billed_for_every_round( + first=response, + first_breakdown=first_round_breakdown, + final=LiteLLM_Proxy_MCP_Handler.add_mcp_output_elements_to_response( + response=final_response, + mcp_tools_fetched=original_mcp_tools, + tool_results=tool_results, + ), + final_breakdown=final_round_breakdown, + ) + return final_response return response @@ -442,6 +540,18 @@ def _bridges_to_chat_completions( _RESPONSES_ONLY_REQUEST_FIELDS_NEVER_BRIDGED: Final = frozenset({"client_metadata"}) +def _previous_response_client_ids(previous_response_id: str | None, kwargs: Mapping[str, object]) -> tuple[str, ...]: + """The forms of the previous response's id the client held, which its spend log row can be stored under.""" + if not previous_response_id: + return () + addressed_id: Final = kwargs.get(ADDRESSED_RESPONSE_ID_FIELD) + return ( + (previous_response_id, addressed_id) + if isinstance(addressed_id, str) and addressed_id + else (previous_response_id,) + ) + + def _bridge_kwargs( kwargs: Mapping[str, object], responses_api_provider_config: BaseResponsesAPIConfig | None, @@ -569,6 +679,18 @@ def _will_bridge_to_chat_completions( ) +def _provider_stores_the_first_round(model: str, custom_llm_provider: str | None, kwargs: Mapping[str, object]) -> bool: + """A bridged first MCP round runs as an internal call, so no spend log holds it for the follow-up to replay.""" + api_base: Final = _api_base_kwarg(kwargs) + return not _will_bridge_to_chat_completions( + model, + custom_llm_provider or litellm.get_llm_provider(model=model, api_base=api_base)[1], + bool(kwargs.get("use_chat_completions_api")), + kwargs.get("model_info"), + api_base, + ) + + def _is_a_client_input_item( item: object, ) -> TypeIs[Mapping[str, object]]: # guard-ok: bounds a client input item before its mark is read @@ -1483,6 +1605,7 @@ def responses( if _bridges_to_chat_completions(responses_api_provider_config, use_chat_completions_api): bridge_kwargs: Final = _bridge_kwargs(kwargs, responses_api_provider_config, allowed_openai_params) return litellm_completion_transformation_handler.response_api_handler( + _previous_response_client_ids(previous_response_id, kwargs), model=model, input=input, responses_api_request=response_api_optional_params, diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index b711d051055..b4dbf4c3d4e 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -258,6 +258,7 @@ class LiteLLM_Proxy_MCP_Handler: request_tags: list[str] | None = None, raw_headers: dict[str, str] | None = None, granted_toolsets: "GrantedToolsetIds | None" = None, + request_trace_id: str | None = None, ) -> tuple[list[MCPTool], list[str]]: """ Get available tools from the MCP server manager. @@ -343,6 +344,7 @@ class LiteLLM_Proxy_MCP_Handler: log_list_tools_to_spendlogs=True, list_tools_log_source="responses", litellm_trace_id=litellm_trace_id, + request_trace_id=request_trace_id, request_tags=request_tags, raw_headers=raw_headers, ) @@ -474,6 +476,7 @@ class LiteLLM_Proxy_MCP_Handler: mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, request_tags: list[str] | None = None, raw_headers: dict[str, str] | None = None, + request_trace_id: str | None = None, ) -> tuple[list[MCPTool], dict[str, str]]: """ Process MCP tools through filtering and deduplication pipeline without OpenAI transformation. @@ -505,6 +508,7 @@ class LiteLLM_Proxy_MCP_Handler: mcp_server_auth_headers=mcp_server_auth_headers, request_tags=request_tags, raw_headers=raw_headers, + request_trace_id=request_trace_id, ) # Step 2: Filter tools based on allowed_tools parameter @@ -729,6 +733,7 @@ class LiteLLM_Proxy_MCP_Handler: request_tags: list[str] | None = None, guardrail_context: Mapping[str, object] | None = None, served_tools: Sequence[MCPTool] | None = None, + request_trace_id: str | None = None, ) -> list[MCPToolResult]: """Execute tool calls and return results.""" from fastapi import HTTPException @@ -787,7 +792,11 @@ class LiteLLM_Proxy_MCP_Handler: }, } ] - tool_logging_call_id = litellm_call_id or str(uuid.uuid4()) + tool_logging_call_id = ( + f"{litellm_call_id}_{tool_call_id or uuid.uuid4()}" + if request_trace_id and litellm_call_id + else litellm_call_id or str(uuid.uuid4()) + ) logging_metadata: dict[str, object] = { "tool_call_id": tool_call_id, "tool_name": sanitized_tool_name, @@ -813,6 +822,8 @@ class LiteLLM_Proxy_MCP_Handler: } if litellm_trace_id: logging_request_data["litellm_trace_id"] = litellm_trace_id + if request_trace_id: + logging_metadata["existing_trace_id"] = request_trace_id if request_tags: logging_metadata["tags"] = request_tags if typed_user_api_key_auth is not None: diff --git a/litellm/responses/mcp/round_billing.py b/litellm/responses/mcp/round_billing.py new file mode 100644 index 00000000000..f20d01785a5 --- /dev/null +++ b/litellm/responses/mcp/round_billing.py @@ -0,0 +1,134 @@ +from asyncio import CancelledError +from collections.abc import Generator, Mapping +from contextlib import contextmanager +from types import MappingProxyType +from typing import Final + +from pydantic import TypeAdapter + +from litellm.litellm_core_utils.hidden_params import served_from_cache +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.responses.utils import ResponseAPILoggingUtils +from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse +from litellm.types.utils import CostBreakdown + +_PER_REQUEST_KEYS: Final = frozenset({"discount_percent", "margin_percent", "margin_fixed_amount"}) +_TOTALS_WITH_FIXED_MARGIN: Final = ("total_cost", "margin_total_amount") +_COST_BREAKDOWN: Final = TypeAdapter(CostBreakdown) +_STR_KEYED: Final = TypeAdapter(Mapping[str, object]) + + +def _summed(first: object, final: object) -> object: + if first is None: + return final + if final is None: + return first + if isinstance(first, bool) or isinstance(final, bool): + return final + if isinstance(first, (int, float)) and isinstance(final, (int, float)): + return first + final + if isinstance(first, Mapping) and isinstance(final, Mapping): + return _summed_mapping(_STR_KEYED.validate_python(first), _STR_KEYED.validate_python(final)) + return final + + +def _summed_entry(key: str, first: Mapping[str, object], final: Mapping[str, object]) -> object: + if key in _PER_REQUEST_KEYS: + return final[key] if key in final else first[key] + return _summed(first.get(key), final.get(key)) + + +def _summed_mapping(first: Mapping[str, object], final: Mapping[str, object]) -> Mapping[str, object]: + keys: Final = (*first, *(key for key in final if key not in first)) + return MappingProxyType({key: _summed_entry(key, first, final) for key in keys}) + + +def merged_round_usage(first: ResponseAPIUsage | None, final: ResponseAPIUsage | None) -> ResponseAPIUsage | None: + if first is None or final is None: + return final if first is None else first + return ResponseAPIUsage.model_validate(_summed_mapping(first.model_dump(), final.model_dump())) + + +def _repeated_fixed_margin( + first: ResponsesAPIResponse, + first_breakdown: CostBreakdown | None, + final: ResponsesAPIResponse, + final_breakdown: CostBreakdown | None, +) -> float: + if first_breakdown is None or final_breakdown is None or first_breakdown is final_breakdown: + return 0.0 + if served_from_cache(first) or served_from_cache(final): + return 0.0 + return min(first_breakdown.get("margin_fixed_amount", 0.0), final_breakdown.get("margin_fixed_amount", 0.0)) + + +def _less_repeated_fixed_margin(summed: Mapping[str, object], repeated: float) -> Mapping[str, object]: + charged_once: Final = MappingProxyType( + {key: total - repeated for key in _TOTALS_WITH_FIXED_MARGIN if isinstance(total := summed.get(key), float)} + ) + return MappingProxyType({**summed, **charged_once}) + + +def summed_cost_breakdown( + first: ResponsesAPIResponse, + first_breakdown: CostBreakdown | None, + final: ResponsesAPIResponse, + final_breakdown: CostBreakdown | None, +) -> CostBreakdown | None: + """The breakdown of the charged rounds, or None when a charged round was priced without its own breakdown, since a + partial one would disagree with the billed cost. The fixed margin is charged once per request, not once per round.""" + if served_from_cache(final): + return None if served_from_cache(first) else first_breakdown + if first_breakdown is final_breakdown: + return None + if served_from_cache(first): + return final_breakdown + if first_breakdown is None or final_breakdown is None: + return None + return _COST_BREAKDOWN.validate_python( + _less_repeated_fixed_margin( + _summed_mapping(first_breakdown, final_breakdown), + _repeated_fixed_margin(first, first_breakdown, final, final_breakdown), + ) + ) + + +def billed_for_every_round( + first: ResponsesAPIResponse, + first_breakdown: CostBreakdown | None, + final: ResponsesAPIResponse, + final_breakdown: CostBreakdown | None, +) -> ResponsesAPIResponse: + """The final round's response carrying the usage and cost of both rounds, so the one logged row bills the request. + + Costs are summed per round rather than repriced from the summed usage, because a round's price can depend on its + own size (tiered rates) or on charges the usage does not show (built-in tools, provider-reported costs). The fixed + margin both rounds added is charged once, since it is a per-request fee. + """ + first_cost: Final = first.hidden_params.get("response_cost") + final_cost: Final = final.hidden_params.get("response_cost") + total_cost: Final = ( + first_cost + final_cost - _repeated_fixed_margin(first, first_breakdown, final, final_breakdown) + if isinstance(first_cost, (int, float)) and isinstance(final_cost, (int, float)) + else None + ) + billed: Final = final.model_copy(update=MappingProxyType({"usage": merged_round_usage(first.usage, final.usage)})) + billed.hidden_params = { + **final.hidden_params, + "response_cost": total_cost, + } + return billed + + +@contextmanager +def first_round_billed_on_failure(logging_obj: object, first: ResponsesAPIResponse) -> Generator[None]: + try: + yield + except (Exception, CancelledError): + first_cost: Final = first.hidden_params.get("response_cost") + if isinstance(logging_obj, LiteLLMLoggingObj) and first.usage is not None and not served_from_cache(first): + logging_obj.record_partial_usage_for_failure( + ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(first.usage), + float(first_cost) if isinstance(first_cost, (int, float)) else 0.0, + ) + raise diff --git a/litellm/utils.py b/litellm/utils.py index 27669d8fd4a..8aa6eb23c68 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1374,15 +1374,16 @@ def _dispatch_success_logging( is_completion_with_fallbacks: bool, is_litellm_internal_call: bool, ) -> None: - if not is_litellm_internal_call: - _schedule_async_success_logging( - logging_obj=logging_obj, - result=result, - start_time=start_time, - end_time=end_time, - is_completion_with_fallbacks=is_completion_with_fallbacks, - ) + if is_litellm_internal_call: + return + _schedule_async_success_logging( + logging_obj=logging_obj, + result=result, + start_time=start_time, + end_time=end_time, + is_completion_with_fallbacks=is_completion_with_fallbacks, + ) logging_obj.handle_sync_success_callbacks_for_async_calls( result=result, start_time=start_time, diff --git a/tests/integration/mcp/test_mcp_llm_endpoints.py b/tests/integration/mcp/test_mcp_llm_endpoints.py index c3bccec1f8e..e288cef3699 100644 --- a/tests/integration/mcp/test_mcp_llm_endpoints.py +++ b/tests/integration/mcp/test_mcp_llm_endpoints.py @@ -47,10 +47,20 @@ pytestmark: Final = pytest.mark.xdist_group("mcp_llm_private_gateway") @pytest.fixture(scope="session") def _llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, Mapping[str, str]]]: directory: Final = tmp_path_factory.mktemp("mcp-llm-gateway") + stock: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + config: Final = directory / "stored-responses.yaml" + config.write_text( + yaml.safe_dump( + { + **stock, + "general_settings": {**object_value(stock["general_settings"]), "store_prompts_in_spend_logs": True}, + } + ) + ) with ( worker_services(directory) as services, gateway_from_environment() as shared, - owned_proxy(shared, directory, services) as private, + owned_proxy(shared, directory, services, config=config) as private, ): yield private, services @@ -106,6 +116,17 @@ class Turn: answer: str +@dataclass(frozen=True, slots=True) +class Billing: + first: tuple[int, int] = (10, 5) + follow_up: tuple[int, int] = (10, 5) + follow_up_fails: bool = False + + +EVEN_ROUNDS: Final = Billing() +FOLLOW_UP_OUTAGE: Final = "scripted follow-up outage" + + def _fixed_turn(tool: str) -> Callable[[Mapping[str, JsonValue]], Turn]: return lambda _: Turn(tool, json.dumps(ADD), ANSWER) @@ -162,7 +183,58 @@ def _responses_stream(response: Mapping[str, JsonValue], item: Mapping[str, Json ) -def _model_double(plan: Callable[[Mapping[str, JsonValue]], Turn]) -> Callable[[Request], Reply]: +def _block_ids(message: JsonValue, kind: str, field: str) -> frozenset[str]: + content: Final = message.get("content") if isinstance(message, dict) else None + blocks: Final = content if isinstance(content, list) else [] + return frozenset(str(block[field]) for block in blocks if isinstance(block, dict) and block.get("type") == kind) + + +def _unanswered_tool_uses(body: Mapping[str, JsonValue]) -> frozenset[str]: + messages: Final = body.get("messages") + turns: Final = messages if isinstance(messages, list) else [] + followers: Final = [*turns[1:], None] + return frozenset().union( + *( + _block_ids(turn, "tool_use", "id") - _block_ids(follower, "tool_result", "tool_use_id") + for turn, follower in zip(turns, followers) + ) + ) + + +def _orphan_tool_results(body: Mapping[str, JsonValue]) -> frozenset[str]: + messages: Final = body.get("messages") + turns: Final = messages if isinstance(messages, list) else [] + leaders: Final = [None, *turns[:-1]] + return frozenset().union( + *( + _block_ids(turn, "tool_result", "tool_use_id") - _block_ids(leader, "tool_use", "id") + for turn, leader in zip(turns, leaders) + ) + ) + + +def _anthropic_invalid_request(message: str) -> Reply: + error: Final = {"type": "error", "error": {"type": "invalid_request_error", "message": message}} + return Reply(status=400, body=json.dumps(error).encode()) + + +def _anthropic_tool_result_rule(body: Mapping[str, JsonValue]) -> Reply | None: + unanswered: Final = _unanswered_tool_uses(body) + orphans: Final = _orphan_tool_results(body) + if unanswered: + return _anthropic_invalid_request( + f"`tool_use` ids were found without `tool_result` blocks immediately after: {', '.join(sorted(unanswered))}" + ) + if orphans: + return _anthropic_invalid_request( + f"unexpected `tool_use_id` found in `tool_result` blocks: {', '.join(sorted(orphans))}" + ) + return None + + +def _model_double( + plan: Callable[[Mapping[str, JsonValue]], Turn], billing: Billing = EVEN_ROUNDS +) -> Callable[[Request], Reply]: def respond(request: Request) -> Reply: if request.method == "GET" and request.target.endswith("/models"): return _json({"object": "list", "data": []}) @@ -172,8 +244,16 @@ def _model_double(plan: Callable[[Mapping[str, JsonValue]], Turn]) -> Callable[[ return Reply(status=500, body=json.dumps(outage).encode()) turn: Final = plan(body) done: Final = _has_tool_result(body) + if done and billing.follow_up_fails: + failure: Final = {"error": {"message": FOLLOW_UP_OUTAGE, "type": "server_error", "code": None}} + return Reply(status=500, body=json.dumps(failure).encode()) + prompt_tokens, completion_tokens = billing.follow_up if done else billing.first identity: Final = uuid.uuid4().hex[:12] - usage: Final = {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15} + usage: Final = { + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": prompt_tokens + completion_tokens, + } if request.target.endswith("/chat/completions"): message: Final = ( {"role": "assistant", "content": turn.answer} @@ -222,6 +302,8 @@ def _model_double(plan: Callable[[Mapping[str, JsonValue]], Turn]) -> Callable[[ } ) if request.target.endswith("/messages"): + if (rejected := _anthropic_tool_result_rule(body)) is not None: + return rejected content: Final = ( [{"type": "text", "text": turn.answer}] if done @@ -236,7 +318,7 @@ def _model_double(plan: Callable[[Mapping[str, JsonValue]], Turn]) -> Callable[[ "content": content, "stop_reason": "end_turn" if done else "tool_use", "stop_sequence": None, - "usage": {"input_tokens": 10, "output_tokens": 5}, + "usage": {"input_tokens": prompt_tokens, "output_tokens": completion_tokens}, } ) assert request.target.endswith("/responses"), request.target @@ -265,7 +347,11 @@ def _model_double(plan: Callable[[Mapping[str, JsonValue]], Turn]) -> Callable[[ "status": "completed", "model": "gpt-4o-mini", "output": [item], - "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + "usage": { + "input_tokens": prompt_tokens, + "output_tokens": completion_tokens, + "total_tokens": prompt_tokens + completion_tokens, + }, } return _responses_stream(response, item) if body.get("stream") is True else _json(response) @@ -354,16 +440,28 @@ def _upstream_model(surface: Surface) -> str: }[surface] +PRICES: Final[Mapping[str, JsonValue]] = {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002, "num_retries": 0} + + @contextmanager -def _rig(gateway: Gateway, surface: Surface) -> Iterator[Rig]: +def _rig( + gateway: Gateway, + surface: Surface, + *, + upstream: Surface | None = None, + billing: Billing = EVEN_ROUNDS, + prices: Mapping[str, JsonValue] | None = None, +) -> Iterator[Rig]: alias: Final = "llm" + uuid.uuid4().hex[:8] with ( mcp_peer() as peer, - wire_server(_model_double(_fixed_turn(f"{alias}-add"))) as wire, + wire_server(_model_double(_fixed_turn(f"{alias}-add"), billing)) as wire, gateway.scenario() as scenario, ): server_id: Final = register_mcp(scenario, peer, alias) - model: Final = scenario.model(model=_upstream_model(surface), api_base=wire.url + "/v1") + model: Final = scenario.model( + model=_upstream_model(upstream or surface), api_base=wire.url + "/v1", **(prices or {}) + ) peer.drain() yield Rig(gateway, scenario, peer, wire, alias, server_id, model, surface) @@ -590,6 +688,119 @@ def test_an_open_server_in_another_gateway_cannot_widen_this_gateways_tools(gate assert deleted.json()["info"]["status"] == "deleted", deleted.text +FIRST_ROUND: Final = (11, 3) +FOLLOW_UP_ROUND: Final = (17, 7) + + +def _price(*rounds: tuple[int, int]) -> float: + input_price, output_price = PRICES["input_cost_per_token"], PRICES["output_cost_per_token"] + assert isinstance(input_price, float) and isinstance(output_price, float) + return sum(prompt * input_price + completion * output_price for prompt, completion in rounds) + + +def _responses_rows(call_id: str, status: str) -> list[dict[str, JsonValue]]: + return eventually( + lambda: read_rows( + "SELECT request_id, spend, prompt_tokens, completion_tokens, session_id " + 'FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s AND call_type=%s AND status=%s', + (call_id, "aresponses", status), + ), + lambda found: len(found) >= 1, + seconds=70, + ) + + +def _session_call_types(session_id: str) -> frozenset[str]: + rows: Final = read_rows('SELECT call_type FROM "LiteLLM_SpendLogs" WHERE session_id=%s', (session_id,)) + return frozenset(str(row["call_type"]) for row in rows) + + +def _message_text(message: Mapping[str, JsonValue]) -> str: + content: Final = message.get("content") + if isinstance(content, str): + return content + blocks: Final = content if isinstance(content, list) else [] + return "".join(str(block["text"]) for block in blocks if isinstance(block, dict) and block.get("type") == "text") + + +def _assistant_texts(body: Mapping[str, JsonValue]) -> tuple[str, ...]: + messages: Final = body.get("messages") + turns: Final = messages if isinstance(messages, list) else [] + return tuple(_message_text(turn) for turn in turns if isinstance(turn, dict) and turn.get("role") == "assistant") + + +def test_nonstream_responses_auto_execute_logs_the_answer_and_bills_both_rounds(gateway: Gateway) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND) + with _rig(gateway, "responses", billing=billing, prices=PRICES) as rig: + response: Final = rig.send(_granted_key(rig), [AUTO]) + assert response.status_code == 200, response.text + body: Final = JSON_OBJECT.validate_json(response.content) + assert rig.final_text(body) == ANSWER, response.text + rows: Final = _responses_rows(response.headers["x-litellm-call-id"], "success") + assert len(rows) == 1, rows + row: Final = rows[0] + assert row["request_id"] == body["id"], row + assert (row["prompt_tokens"], row["completion_tokens"]) == ( + FIRST_ROUND[0] + FOLLOW_UP_ROUND[0], + FIRST_ROUND[1] + FOLLOW_UP_ROUND[1], + ), row + assert object_value(body["usage"])["input_tokens"] == row["prompt_tokens"], body + assert float(str(row["spend"])) == pytest.approx(_price(FIRST_ROUND, FOLLOW_UP_ROUND)), row + assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(float(str(row["spend"]))) + session_id: Final = string_value(row["session_id"]) + eventually( + lambda: _session_call_types(session_id), + lambda found: {"list_mcp_tools", "call_mcp_tool"} <= found, + seconds=70, + ) + + +def test_nonstream_responses_failed_follow_up_bills_round_one_on_the_failure_row(gateway: Gateway) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, follow_up_fails=True) + with _rig(gateway, "responses", billing=billing, prices=PRICES) as rig: + response: Final = rig.send(_granted_key(rig), [AUTO]) + assert response.status_code == 500 and FOLLOW_UP_OUTAGE in response.text, response.text + call_id: Final = response.headers["x-litellm-call-id"] + rows: Final = _responses_rows(call_id, "failure") + assert len(rows) == 1, rows + assert float(str(rows[0]["spend"])) == pytest.approx(_price(FIRST_ROUND)), rows + assert ( + read_rows( + 'SELECT request_id FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s AND call_type=%s AND status=%s', + (call_id, "aresponses", "success"), + ) + == [] + ) + + +def test_bridged_responses_auto_execute_answers_every_tool_use_with_its_result(gateway: Gateway) -> None: + with _rig(gateway, "responses", upstream="messages") as rig: + response: Final = rig.send(_granted_key(rig), [AUTO]) + assert response.status_code == 200, response.text + assert rig.final_text(response.json()) == ANSWER, response.text + assert [call["body"]["params"]["name"] for call in _peer_add_calls(rig.peer)] == ["add"] + + +def test_bridged_next_turn_replays_the_mcp_turn_answer(gateway: Gateway) -> None: + with _rig(gateway, "responses", upstream="messages") as rig: + key: Final = _granted_key(rig) + first: Final = rig.send(key, [AUTO]) + assert first.status_code == 200, first.text + first_id: Final = string_value(JSON_OBJECT.validate_json(first.content)["id"]) + eventually( + lambda: read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (first_id,)), + lambda found: len(found) == 1, + seconds=70, + ) + rig.wire.drain() + second: Final = rig.send(key, [AUTO], previous_response_id=first_id, cache={"no-cache": True}) + assert second.status_code == 200, second.text + opening: Final = next( + JSON_OBJECT.validate_json(request.body) for request in rig.wire.drain() if request.method == "POST" + ) + assert ANSWER in _assistant_texts(opening), opening + + Bridge = Literal["chat", "responses", "messages"] BRIDGES: Final[tuple[Bridge, ...]] = ("chat", "responses") Client = Literal["sync", "async"] diff --git a/tests/unit/caching/test_caching_handler.py b/tests/unit/caching/test_caching_handler.py index a9541f9fcb8..58549a028ea 100644 --- a/tests/unit/caching/test_caching_handler.py +++ b/tests/unit/caching/test_caching_handler.py @@ -52,7 +52,8 @@ import json import httpx import respx from fastapi.testclient import TestClient -from litellm._internal_context import current_service_target, in_post_response_phase +from contextlib import nullcontext +from litellm._internal_context import current_service_target, in_post_response_phase, internal_sub_call from litellm.caching.caching_handler import _PENDING_CACHE_WRITES from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler @@ -3046,3 +3047,57 @@ async def test_async_get_cache_answers_400_for_a_single_object_embedding_input() call_type=CallTypes.aembedding.value, kwargs={"model": model, "custom_llm_provider": "gemini", "input": clip_block, "caching": True}, ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("internal", [False, True]) +async def test_a_cache_hit_inside_an_internal_sub_call_is_free_and_fires_no_callbacks_of_its_own( + monkeypatch: pytest.MonkeyPatch, internal: bool +): + cache: Final = Cache(type="local") + monkeypatch.setattr(litellm, "cache", cache) + kwargs: Final = {"model": "gpt-5.4", "messages": [{"role": "user", "content": "hello"}], "caching": True} + cache.add_cache( + result=litellm.ModelResponse( + model="gpt-5.4", + choices=[litellm.Choices(message=litellm.Message(role="assistant", content="hi"))], + usage=litellm.Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500), + ), + **kwargs, + ) + async_logging: Final = _build_logging_obj(CallTypes.acompletion.value, stream=False) + async_logging.async_success_handler = AsyncMock() + async_logging.handle_sync_success_callbacks_for_async_calls = MagicMock() + sync_logging: Final = _build_logging_obj(CallTypes.completion.value, stream=False) + sync_logging.handle_sync_success_callbacks_for_async_calls = MagicMock() + + with internal_sub_call() if internal else nullcontext(): + async_hit: Final = await LLMCachingHandler( + original_function=litellm.acompletion, request_kwargs=kwargs, start_time=_FIXED_START + ).async_get_cache( + model="gpt-5.4", + original_function=litellm.acompletion, + logging_obj=async_logging, + start_time=_FIXED_START, + call_type=CallTypes.acompletion.value, + kwargs=kwargs, + args=(), + ) + sync_hit: Final = LLMCachingHandler( + original_function=completion, request_kwargs=kwargs, start_time=_FIXED_START + ).sync_get_cache( + model="gpt-5.4", + original_function=completion, + logging_obj=sync_logging, + start_time=_FIXED_START, + call_type=CallTypes.completion.value, + kwargs=kwargs, + ) + + assert async_hit is not None and async_hit.cached_result is not None + assert sync_hit.cached_result is not None + assert async_hit.cached_result._hidden_params.get("response_cost") == 0.0 + assert sync_hit.cached_result._hidden_params.get("response_cost") == 0.0 + assert async_logging.async_success_handler.called is not internal + assert async_logging.handle_sync_success_callbacks_for_async_calls.called is not internal + assert sync_logging.handle_sync_success_callbacks_for_async_calls.called is not internal diff --git a/tests/unit/litellm_core_utils/test_hidden_params.py b/tests/unit/litellm_core_utils/test_hidden_params.py index 9d1e1e8dceb..53b6c0370fa 100644 --- a/tests/unit/litellm_core_utils/test_hidden_params.py +++ b/tests/unit/litellm_core_utils/test_hidden_params.py @@ -7,6 +7,7 @@ from litellm.litellm_core_utils.hidden_params import ( get_hidden_params, get_hidden_params_storage, get_or_create_hidden_params, + served_from_cache, set_hidden_param, set_hidden_params, ) @@ -250,3 +251,24 @@ def test_set_hidden_params_replaces_frozen_decisions_response_private_attr() -> assert response.hidden_params is replacement assert response._hidden_params is replacement + + +@pytest.mark.parametrize( + "hidden_params, expected", + [ + ({"cache_hit": True}, True), + ({"cache_hit": None}, False), + ({"cache_hit": False}, False), + ({"cache_hit": "True"}, False), + ({}, False), + ], +) +def test_served_from_cache_is_true_only_for_a_response_the_cache_stamped( + hidden_params: dict[str, object], expected: bool +) -> None: + model_response: Final = ModelResponse() + set_hidden_params(model_response, dict(hidden_params)) + + assert served_from_cache(model_response) is expected + assert served_from_cache({HIDDEN_PARAMS_ATTR: dict(hidden_params)}) is expected + assert served_from_cache(object()) is False diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index df6593d29b1..7f2653cb481 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -823,6 +823,20 @@ class TestZeroCostDiagnostic: assert logging_obj.model_call_details["standard_logging_object"]["zero_cost_diagnostic"] is None assert self._zero_cost_warnings(caplog) == [] + def test_response_served_from_the_cache_at_no_charge_stays_silent( + self, deployment_pricing: Mapping[str, float], caplog: pytest.LogCaptureFixture + ) -> None: + usage: Final = litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30) + logging_obj: Final = self._logging_obj(deployment_pricing) + response: Final = self._response(usage, response_cost=0.0, cache_hit=True, model_id=self.DEPLOYMENT_ID) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + cost: Final = logging_obj.response_cost_calculator(result=response) + + assert cost == 0.0 + assert logging_obj.model_call_details["zero_cost_diagnostic"] is None + assert self._zero_cost_warnings(caplog) == [] + def test_unbilled_read_route_with_usage_stays_silent( self, deployment_pricing: Mapping[str, float], caplog: pytest.LogCaptureFixture ) -> None: @@ -5363,6 +5377,63 @@ def test_get_standard_logging_object_payload_includes_litellm_call_id(logging_ob assert payload["litellm_call_id"] == call_id +@pytest.mark.parametrize( + "hidden_params, passed_cache_hit, expected_cache_hit, expected_cost, expected_unique_id", + [ + ({"cache_hit": True, "response_cost": 0.0}, None, True, 0.0, True), + ({"cache_hit": True, "response_cost": 0.002}, None, None, 0.002, True), + ({"cache_hit": True, "response_cost": 0.0}, False, False, 0.0, True), + ({"response_cost": 0.0}, None, None, 0.0, False), + ({}, None, None, None, False), + ], +) +def test_success_handler_logs_a_response_served_from_the_cache_under_a_unique_id( + hidden_params: dict[str, object], + passed_cache_hit: bool | None, + expected_cache_hit: bool | None, + expected_cost: float | None, + expected_unique_id: bool, +): + """A parent call whose response came back from the response cache (an internal sub-call's hit) reuses the + id of the row that cached it, so its row needs a unique id, and it is a cache hit only when it cost nothing.""" + from litellm.litellm_core_utils.hidden_params import set_hidden_params + + logging_obj: Final = LitellmLogging( + model="gpt-5.6", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="acompletion", + start_time=time.time(), + litellm_call_id="parent-call-id", + function_id="parent-fn-id", + ) + logging_obj.update_environment_variables( + litellm_params={}, optional_params={}, model="gpt-5.6", custom_llm_provider="openai" + ) + response: Final = ModelResponse( + id="chatcmpl-cached-src", + model="gpt-5.6", + usage=Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500), + ) + set_hidden_params(response, dict(hidden_params)) + + logging_obj._success_handler_helper_fn( + result=response, + start_time=datetime.datetime.now(), + end_time=datetime.datetime.now(), + cache_hit=passed_cache_hit, + ) + + payload: Final = logging_obj.model_call_details["standard_logging_object"] + assert logging_obj.model_call_details["cache_hit"] is expected_cache_hit + assert payload["cache_hit"] is expected_cache_hit + if expected_cost is not None: + assert payload["response_cost"] == expected_cost + assert payload["id"].startswith("chatcmpl-cached-src_cache_hit") is expected_unique_id + assert (payload["id"] == "chatcmpl-cached-src") is not expected_unique_id + assert payload["saved_cache_cost"] == 0.0 or expected_cache_hit is True + + @pytest.mark.parametrize( "client_sent_oauth_token, custom_llm_provider, expected", [(True, "anthropic", True), (True, "bedrock", False), (False, "anthropic", False), (None, "anthropic", None)], diff --git a/tests/unit/proxy/db/test_spend_log_tool_index.py b/tests/unit/proxy/db/test_spend_log_tool_index.py index eaa94c0438f..af3268a7cc9 100644 --- a/tests/unit/proxy/db/test_spend_log_tool_index.py +++ b/tests/unit/proxy/db/test_spend_log_tool_index.py @@ -21,6 +21,8 @@ from litellm.proxy.db.spend_log_tool_index import ( flush_tool_usage_transactions, response_tool_call_names, ) +from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler +from litellm.types.llms.openai import ResponsesAPIResponse from tests.unit.proxy.db.fake_prisma_engine import engine_call @@ -222,6 +224,37 @@ class TestResponseToolCallNames: } assert response_tool_call_names(response) == ("ant_get_weather",) + def test_auto_executed_mcp_tools_count_from_the_final_responses_api_response(self): + final_round = ResponsesAPIResponse.model_validate( + { + "id": "resp_final", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-5.6", + "output": [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "done", "annotations": []}], + } + ], + } + ) + logged = LiteLLM_Proxy_MCP_Handler.add_mcp_output_elements_to_response( + final_round, + mcp_tools_fetched=[{"name": "demo-echo"}, {"name": "demo-add"}], + tool_results=[ + {"tool_call_id": "call_1", "result": "x", "name": "demo-echo"}, + {"tool_call_id": "call_2", "result": "3", "name": "demo-add"}, + ], + ) + + assert response_tool_call_names(logged) == ("demo-echo", "demo-add") + assert response_tool_call_names(logged.model_dump()) == ("demo-echo", "demo-add") + def _transaction( request_id: str, diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index de2676bafaa..65dcce97b90 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -4106,6 +4106,36 @@ def test_get_logging_payload_cache_hit_keeps_raw_litellm_call_id(): assert json.loads(payload["metadata"])["litellm_call_id"] != payload["request_id"] +@pytest.mark.parametrize("cached", [True, False]) +def test_get_logging_payload_gives_a_response_served_from_the_cache_a_unique_request_id(cached: bool): + """A parent call that returns an internal sub-call's cached response reuses the cached response's id, + and SpendLogs skips a duplicate request_id, so the row would be dropped unless its id is made unique.""" + from litellm.litellm_core_utils.hidden_params import set_hidden_params + + response_obj: Final = litellm.ModelResponse( + id="chatcmpl-cached-src", + model="gpt-5.6", + usage=litellm.Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2), + ) + set_hidden_params(response_obj, {"cache_hit": True, "response_cost": 0.0} if cached else {}) + now: Final = datetime.datetime.now(timezone.utc) + + payload: Final = get_logging_payload( + kwargs={ + "model": "openai/gpt-5.6", + "call_type": "aresponses", + "litellm_call_id": "parent-call-id", + "litellm_params": {"metadata": {"user_api_key": "sk-test"}}, + }, + response_obj=response_obj, + start_time=now, + end_time=now, + ) + + assert payload["request_id"].startswith("chatcmpl-cached-src_cache_hit") is cached + assert (payload["request_id"] == "chatcmpl-cached-src") is not cached + + class TestSpendLogKeyRedaction: """Regression: plaintext API keys with Bearer prefix were stored in SpendLogs for failed requests (LIT-4121)""" diff --git a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py index 61b23428d62..025f1315aab 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py +++ b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py @@ -50,6 +50,7 @@ async def test_async_response_api_handler_merges_trace_id_without_error(): litellm_completion_request: dict[str, object], instructions: str | None = None, request_kwargs: Mapping[str, object] | None = None, + previous_response_client_ids: tuple[str, ...] = (), ) -> dict[str, object]: litellm_completion_request["litellm_trace_id"] = "session-trace" return litellm_completion_request diff --git a/tests/unit/responses/litellm_completion_transformation/test_handler.py b/tests/unit/responses/litellm_completion_transformation/test_handler.py index bb374b90f4e..d731bf5f41c 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_handler.py +++ b/tests/unit/responses/litellm_completion_transformation/test_handler.py @@ -13,7 +13,8 @@ kwargs lack the flag and these tests fail. """ import json -from collections.abc import Mapping +from collections.abc import Mapping, Sequence +from types import SimpleNamespace from typing import Final from unittest.mock import patch @@ -25,6 +26,7 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.responses.litellm_completion_transformation.handler import ( LiteLLMCompletionTransformationHandler, ) +from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.utils import ADDRESSED_RESPONSE_ID_FIELD @@ -223,3 +225,88 @@ async def test_bridged_follow_up_turn_keeps_the_addressed_response_id_off_the_pr ) assert isinstance(response, ResponsesAPIResponse) assert [item.type for item in response.output] == ["message"] + + +@pytest.mark.asyncio +async def test_bridged_turn_with_a_client_field_named_previous_response_client_ids_reaches_the_provider(): + provider: Final = _RecordingAnthropicHandler(_ANTHROPIC_MESSAGE_PAYLOAD) + client: Final = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(provider)) + + response = await litellm.aresponses( + model="azure_ai/claude-sonnet-4-6", + api_base="https://fake-foundry-resource.services.ai.azure.com", + api_key="fake-api-key", + input="Double it", + client=client, + previous_response_client_ids=["resp_from_the_client_body"], + ) + + assert provider.request_body is not None, "the bridged turn never reached the provider" + assert isinstance(response, ResponsesAPIResponse) + + +class _StoredTurns: + def __init__(self, rows: tuple[Mapping[str, object], ...]) -> None: + self.rows: Final = rows + + async def query_raw(self, query: str, request_ids: Sequence[str]) -> list[dict[str, object]]: + sessions: Final = {row["session_id"] for row in self.rows if row["request_id"] in request_ids} + return [dict(row) for row in self.rows if row["session_id"] in sessions] + + +def _stored_mcp_turn(request_id: str) -> Mapping[str, object]: + return { + "request_id": request_id, + "call_type": "aresponses", + "session_id": "turn-one-session", + "proxy_server_request": {"input": "My favorite number is 7.", "model": "claude-sonnet-4-6"}, + "response": { + "id": request_id, + "object": "response", + "output": [ + { + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "Noted.", "annotations": []}], + }, + {"type": "mcp_tools_fetched", "tools": [{"name": "demo-echo"}]}, + ], + }, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stored_under", ["litellm_encoded_id", "proxy_issued_id"]) +async def test_bridged_next_turn_replays_a_turn_stored_under_the_id_the_client_holds(stored_under: str): + litellm_encoded_id: Final = ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="azure_ai", model_id="deployment-a", response_id="chatcmpl-turn-one" + ) + proxy_issued_id: Final = "resp_proxy-issued-turn-one" + stored_turns: Final = _StoredTurns( + (_stored_mcp_turn(litellm_encoded_id if stored_under == "litellm_encoded_id" else proxy_issued_id),) + ) + provider: Final = _RecordingAnthropicHandler(_ANTHROPIC_MESSAGE_PAYLOAD) + client: Final = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(provider)) + + with patch("litellm.proxy.proxy_server.prisma_client", SimpleNamespace(db=stored_turns)): + await litellm.aresponses( + model="azure_ai/claude-sonnet-4-6", + api_base="https://fake-foundry-resource.services.ai.azure.com", + api_key="fake-api-key", + input="Double it", + previous_response_id=litellm_encoded_id, + client=client, + **{ADDRESSED_RESPONSE_ID_FIELD: proxy_issued_id}, + ) + + assert provider.request_body is not None, "the bridged turn never reached the provider" + assert [ + (message["role"], json.dumps(message["content"])) for message in provider.request_body["messages"] + ] == [ + ("user", json.dumps([{"type": "text", "text": "My favorite number is 7."}])), + ("assistant", json.dumps([{"type": "text", "text": "Noted."}])), + ("user", json.dumps([{"type": "text", "text": "Double it"}])), + ] diff --git a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py index 5db563be4b0..7a146c1b665 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py +++ b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py @@ -1,5 +1,7 @@ import json import sqlite3 +from collections.abc import Sequence +from types import SimpleNamespace from typing import Final from unittest.mock import AsyncMock, patch @@ -16,7 +18,9 @@ from litellm.responses.litellm_completion_transformation.session_handler import ResponsesSessionHandler, _normalize_redacted_tool_call_arguments, ) +from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.utils import Message @@ -127,7 +131,7 @@ async def test_get_chat_completion_message_history_for_previous_response_id(): ) # Verify the mock was called with correct parameters - mock_get_spend_logs.assert_called_once_with(previous_response_id) + mock_get_spend_logs.assert_called_once_with(previous_response_id, ()) # Verify the returned ChatCompletionSession structure assert "messages" in result @@ -461,8 +465,8 @@ class _FakePrismaDB: self._results = list(results) self.calls = [] - async def query_raw(self, query, *args): - self.calls.append(args) + async def query_raw(self, query, request_ids): + self.calls.append(tuple(request_ids)) if not self._results: return [] return list(self._results.pop(0)) @@ -555,7 +559,7 @@ async def test_message_history_retries_a_spend_log_the_batch_writer_has_not_flus ("assistant", "OK"), ] assert result["litellm_session_id"] == session_id - assert fake_prisma_client.db.calls == [(request_id,), (request_id,)] + assert [call[0] for call in fake_prisma_client.db.calls] == [request_id, request_id] @pytest.mark.asyncio @@ -623,7 +627,91 @@ async def test_message_history_looks_up_the_decoded_chat_completion_id(): encoded_response_id ) - assert fake_prisma_client.db.calls == [(request_id,)] + assert len(fake_prisma_client.db.calls) == 1 + assert request_id in fake_prisma_client.db.calls[0] + + +class _StoredSpendLogs: + def __init__(self, rows: tuple[dict, ...]) -> None: + self.rows = rows + + async def query_raw(self, query: str, request_ids: Sequence[str]) -> list[dict]: + sessions: Final = {row["session_id"] for row in self.rows if row["request_id"] in request_ids} + return [row for row in self.rows if row["session_id"] in sessions] + + +@pytest.mark.asyncio +async def test_message_history_replays_a_bridged_mcp_turn_stored_under_the_client_response_id(): + client_response_id: Final = ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="anthropic", + model_id="e0f302a1412e78470ebb28cbed01fff5f88c0d331c667e9f2ba4b413c6fbd282", + response_id="chatcmpl-mcp-final-round", + ) + session_id: Final = "mcp-session" + stored: Final = _StoredSpendLogs( + ( + { + "request_id": "list-tools-row", + "call_type": "list_mcp_tools", + "session_id": session_id, + "proxy_server_request": {}, + "response": [{"name": "demo-echo"}], + }, + { + "request_id": client_response_id, + "call_type": "aresponses", + "session_id": session_id, + "proxy_server_request": {"input": "Call echo with PING.", "model": "claude-bridge"}, + "response": { + "id": client_response_id, + "object": "response", + "output": [ + { + "type": "reasoning", + "id": "rs_mcp-final-round", + "summary": [{"type": "summary_text", "text": "The echo tool returned PING."}], + }, + { + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "PING", "annotations": []}], + }, + *LiteLLM_Proxy_MCP_Handler.add_mcp_output_elements_to_response( + ResponsesAPIResponse(id=client_response_id, created_at=0, output=[]), + mcp_tools_fetched=[{"name": "demo-echo"}], + tool_results=[{"tool_name": "demo-echo", "result": "PING"}], + ).output, + ], + }, + }, + { + "request_id": "call-tool-row", + "call_type": "call_mcp_tool", + "session_id": session_id, + "proxy_server_request": {"name": "echo", "arguments": {"value": "PING"}}, + "response": {"content": [{"type": "text", "text": "PING"}], "is_error": False}, + }, + ) + ) + + with patch("litellm.proxy.proxy_server.prisma_client", SimpleNamespace(db=stored)): + result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( + client_response_id + ) + + assert [(message.get("role"), _text_of(message.get("content"))) for message in result["messages"]] == [ + ("user", "Call echo with PING."), + ("assistant", "PING"), + ] + assert result["messages"][-1].get("reasoning_content") == "The echo tool returned PING." + assert result["litellm_session_id"] == session_id + + +def _text_of(content: object) -> object: + if isinstance(content, list): + return "".join(part.get("text", "") for part in content if isinstance(part, dict)) + return content @pytest.mark.asyncio @@ -644,7 +732,7 @@ async def test_session_lookup_does_not_retry_when_spend_logs_are_disabled( ) assert spend_logs == [] - assert fake_prisma_client.db.calls == [("chatcmpl-does-not-exist",)] + assert len(fake_prisma_client.db.calls) == 1 def test_normalize_redacted_arguments_skips_custom_tool_calls(): @@ -801,7 +889,9 @@ class _SqliteBackedPrismaDB: self._connection.executemany('INSERT INTO "LiteLLM_SpendLogs" VALUES (?, ?, ?, ?, ?, ?)', rows) async def query_raw(self, query: str, *args: object) -> list[dict[str, object]]: - return [dict(row) for row in self._connection.execute(query.replace("$1", "?"), args)] + sqlite_query: Final = query.replace("= ANY($1::text[])", "IN (SELECT value FROM json_each(?))").replace("$1", "?") + sqlite_args: Final = tuple(json.dumps(arg) if isinstance(arg, list) else arg for arg in args) + return [dict(row) for row in self._connection.execute(sqlite_query, sqlite_args)] @pytest.mark.asyncio diff --git a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py index 745a48d0f6e..b9cf201a47a 100644 --- a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py +++ b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py @@ -1,18 +1,29 @@ +import asyncio import json import logging -from collections.abc import Mapping, Sequence -from types import MappingProxyType +import sys +from collections.abc import Callable, Iterator, Mapping, Sequence +from contextlib import contextmanager, nullcontext +from datetime import datetime +from types import MappingProxyType, SimpleNamespace +from typing import Any, Final, List, cast +from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest import respx +from mcp.types import CallToolResult, ListToolsResult, TextContent from mcp.types import Tool as MCPTool -from typing import Final, List, Any, cast -from unittest.mock import AsyncMock, patch # Import required modules import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.langfuse.langfuse import langfuse_trace_id, langfuse_trace_metadata +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.proxy._experimental.mcp_server import mcp_server_manager +from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ServerListOk +from litellm.proxy.utils import _failure_fields_to_lift from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler from litellm.types.llms.openai import ( ResponsesAPIResponse, @@ -20,6 +31,7 @@ from litellm.types.llms.openai import ( OpenAIMcpServerTool, ToolParam, ) +from litellm.types.utils import StandardLoggingPayload, Usage class MockUserAPIKeyAuth: @@ -1284,6 +1296,67 @@ async def test_no_duplicate_mcp_tools_in_streaming_e2e(): } +_ROUND_ONE_BODY: Final = MappingProxyType( + { + "id": "resp_round_one", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-5.6", + "output": ( + MappingProxyType( + { + "type": "function_call", + "id": "fc_round_one", + "call_id": "call_echo", + "name": "echo", + "arguments": '{"value": "PING"}', + "status": "completed", + } + ), + ), + "usage": MappingProxyType( + { + "input_tokens": 71, + "input_tokens_details": MappingProxyType({"cached_tokens": 32}), + "output_tokens": 19, + "output_tokens_details": MappingProxyType({"reasoning_tokens": 0}), + "total_tokens": 90, + } + ), + } +) +_ROUND_TWO_BODY: Final = MappingProxyType( + { + "id": "resp_round_two", + "object": "response", + "created_at": 2, + "status": "completed", + "model": "gpt-5.6", + "output": ( + MappingProxyType( + { + "type": "message", + "id": "msg_round_two", + "status": "completed", + "role": "assistant", + "content": (MappingProxyType({"type": "output_text", "text": "PING", "annotations": ()}),), + } + ), + ), + "usage": MappingProxyType( + { + "input_tokens": 149, + "input_tokens_details": MappingProxyType({"cached_tokens": 64}), + "output_tokens": 5, + "output_tokens_details": MappingProxyType({"reasoning_tokens": 0}), + "total_tokens": 154, + } + ), + } +) + + WEATHER_TOOL: Final = MCPTool( name="get_weather", description="Get weather for a city", @@ -1398,6 +1471,769 @@ SCRIPTED_PROVIDERS: Final = MappingProxyType( ) +def _round_cost(body: Mapping[str, object]) -> float: + return litellm.completion_cost( + completion_response=ResponsesAPIResponse.model_validate(body), + model="gpt-5.6", + custom_llm_provider="openai", + call_type="aresponses", + ) + + +class _RoundRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logged: tuple[StandardLoggingPayload, ...] = () + self.failed: tuple[StandardLoggingPayload, ...] = () + self.hooked: tuple[object, ...] = () + + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self.logged = (*self.logged, cast(StandardLoggingPayload, kwargs["standard_logging_object"])) + + async def async_log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self.failed = (*self.failed, cast(StandardLoggingPayload, kwargs["standard_logging_object"])) + + def logging_hook(self, kwargs: dict, result: Any, call_type: str) -> tuple[dict, Any]: + self.hooked = (*self.hooked, result) + return kwargs, result + + +def _no_op_sync_callback( + kwargs: Mapping[str, object], completion_response: object, start_time: datetime, end_time: datetime +) -> None: + return None + + +async def _settle(condition: Callable[[], bool]) -> None: + for _ in range(200): + if condition(): + break + await asyncio.sleep(0.01) + await asyncio.sleep(0.2) + + +_ECHO_MCP_TOOL: Final = MappingProxyType( + {"type": "mcp", "server_url": "litellm_proxy/mcp/demo", "require_approval": "never"} +) + + +def _serve_responses( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, responses: Sequence[httpx.Response] +) -> respx.Route: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + return respx_mock.post(url__regex=r"https://mcp-rounds\.example\.com/.*responses").mock(side_effect=list(responses)) + + +def _serve_rounds( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, bodies: Sequence[Mapping[str, object]] +) -> respx.Route: + return _serve_responses( + respx_mock, monkeypatch, tuple(httpx.Response(200, content=json.dumps(body, default=dict)) for body in bodies) + ) + + +@contextmanager +def _echo_tool_served(tool_result: str) -> Iterator[None]: + echo_tool: Final = MCPTool.model_validate( + {"name": "echo", "inputSchema": {"type": "object", "properties": {"value": {"type": "string"}}}}, + by_name=False, + ) + with ( + patch.object( + LiteLLM_Proxy_MCP_Handler, + "process_mcp_tools_without_openai_transform", + AsyncMock(return_value=([echo_tool], {"echo": "demo"})), + ), + patch.object( + LiteLLM_Proxy_MCP_Handler, + "execute_tool_calls", + AsyncMock(return_value=[{"tool_call_id": "call_echo", "result": tool_result, "name": "echo"}]), + ), + ): + yield + + +async def _auto_execute( + logging_obj: Logging, mcp_tool: Mapping[str, object], tool_result: str +) -> ResponsesAPIResponse: + with _echo_tool_served(tool_result): + response: Final = await litellm.aresponses( + model="openai/gpt-5.6", + input="Call echo with PING", + tools=[dict(mcp_tool)], + api_key="sk-mcp-rounds", + api_base="https://mcp-rounds.example.com/v1", + litellm_logging_obj=logging_obj, + ) + assert isinstance(response, ResponsesAPIResponse) + return response + + +async def _auto_execute_echo( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, logging_obj: Logging +) -> ResponsesAPIResponse: + _serve_rounds(respx_mock, monkeypatch, (_ROUND_ONE_BODY, _ROUND_TWO_BODY)) + return await _auto_execute(logging_obj, _ECHO_MCP_TOOL, "PING") + + +async def _auto_execute_through_mcp_server( + monkeypatch: pytest.MonkeyPatch, call_tool: AsyncMock, logging_obj: Logging +) -> None: + monkeypatch.setitem( + sys.modules, "litellm.proxy.proxy_server", SimpleNamespace(proxy_logging_obj=None, prisma_client=None) + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + SimpleNamespace( + catalog=SimpleNamespace(operation=nullcontext), + get_registry=MagicMock(return_value={}), + call_tool=call_tool, + get_mcp_server_from_tool_name=MagicMock(return_value=None), + get_mcp_server_by_name=MagicMock(return_value=None), + ), + ) + echo_tool: Final = MCPTool.model_validate( + {"name": "echo", "inputSchema": {"type": "object", "properties": {"value": {"type": "string"}}}}, + by_name=False, + ) + with patch.object( + LiteLLM_Proxy_MCP_Handler, + "process_mcp_tools_without_openai_transform", + AsyncMock(return_value=([echo_tool], {"echo": "demo"})), + ): + await litellm.aresponses( + model="openai/gpt-5.6", + input="Call echo with PING", + tools=[dict(_ECHO_MCP_TOOL)], + api_key="sk-mcp-rounds", + api_base="https://mcp-rounds.example.com/v1", + litellm_call_id=logging_obj.litellm_call_id, + litellm_logging_obj=logging_obj, + ) + + +def _rounds_logging_obj(recorder: _RoundRecorder, call_id: str) -> Logging: + return Logging( + model="openai/gpt-5.6", + messages=[{"role": "user", "content": "Call echo with PING"}], + stream=False, + call_type="aresponses", + start_time=datetime(2026, 1, 1), + litellm_call_id=call_id, + function_id=call_id, + dynamic_async_success_callbacks=[recorder], + dynamic_success_callbacks=[recorder, _no_op_sync_callback], + dynamic_async_failure_callbacks=[recorder], + ) + + +def _output_types(output: Sequence[object]) -> tuple[object, ...]: + return tuple(item.get("type") if isinstance(item, Mapping) else getattr(item, "type", None) for item in output) + + +def _assert_one_row_bills_both_rounds(logged: tuple[StandardLoggingPayload, ...], client: ResponsesAPIResponse) -> None: + assert len(logged) == 1 + row: Final = logged[0] + both_rounds: Final = _round_cost(_ROUND_ONE_BODY) + _round_cost(_ROUND_TWO_BODY) + assert row["id"] == client.id + assert _output_types(row["response"]["output"]) == _output_types(client.output) + assert "message" in _output_types(client.output) + assert "function_call" not in _output_types(client.output) + assert (row["prompt_tokens"], row["completion_tokens"]) == (71 + 149, 19 + 5) + assert row["metadata"]["usage_object"]["prompt_tokens_details"]["cached_tokens"] == 32 + 64 + assert row["response_cost"] == pytest.approx(both_rounds) + assert row["cost_breakdown"]["total_cost"] == pytest.approx(both_rounds) + assert client.hidden_params["response_cost"] == pytest.approx(both_rounds) + assert client.usage is not None + assert (client.usage.input_tokens, client.usage.output_tokens) == (71 + 149, 19 + 5) + + +@pytest.mark.asyncio +async def test_non_stream_auto_execute_logs_the_client_response_billed_for_both_rounds( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + recorder: Final = _RoundRecorder() + response: Final = await _auto_execute_echo( + respx_mock, monkeypatch, _rounds_logging_obj(recorder, "mcp-rounds-immediate") + ) + + await _settle(lambda: len(recorder.logged) > 0) + + _assert_one_row_bills_both_rounds(recorder.logged, response) + + +@pytest.mark.asyncio +async def test_non_stream_auto_execute_deferred_for_a_post_call_guardrail_logs_both_rounds_once( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + recorder: Final = _RoundRecorder() + logging_obj: Final = _rounds_logging_obj(recorder, "mcp-rounds-deferred") + logging_obj.defer_async_logging = True + response: Final = await _auto_execute_echo(respx_mock, monkeypatch, logging_obj) + + enqueue_deferred: Final = logging_obj.enqueue_deferred_logging + assert enqueue_deferred is not None + enqueue_deferred() + await _settle(lambda: len(recorder.logged) > 0) + + _assert_one_row_bills_both_rounds(recorder.logged, response) + + +@pytest.mark.asyncio +async def test_non_stream_auto_execute_hands_sync_callbacks_the_client_response( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + recorder: Final = _RoundRecorder() + response: Final = await _auto_execute_echo(respx_mock, monkeypatch, _rounds_logging_obj(recorder, "mcp-rounds-sync")) + + await _settle(lambda: len(recorder.hooked) > 0) + + assert len(recorder.hooked) == 1 + hooked: Final = recorder.hooked[0] + assert isinstance(hooked, ResponsesAPIResponse) + assert hooked.id == response.id + assert _output_types(hooked.output) == _output_types(response.output) + assert hooked.usage is not None + assert (hooked.usage.model_dump().get("prompt_tokens"), hooked.usage.model_dump().get("completion_tokens")) == ( + 71 + 149, + 19 + 5, + ) + + +@pytest.mark.asyncio +async def test_non_stream_auto_execute_bills_the_paid_first_round_when_the_follow_up_fails( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + context_window_error: Final = { + "error": { + "message": "Your input exceeds the context window of this model.", + "type": "invalid_request_error", + "param": "input", + "code": "context_length_exceeded", + } + } + _serve_responses( + respx_mock, + monkeypatch, + ( + httpx.Response(200, content=json.dumps(_ROUND_ONE_BODY, default=dict)), + httpx.Response(400, json=context_window_error), + ), + ) + call_tool: Final = AsyncMock(return_value=CallToolResult(content=[TextContent(type="text", text="PING")])) + recorder: Final = _RoundRecorder() + + with pytest.raises(litellm.BadRequestError): + await _auto_execute_through_mcp_server( + monkeypatch, call_tool, _rounds_logging_obj(recorder, "mcp-rounds-follow-up-fails") + ) + await _settle(lambda: len(recorder.failed) > 0) + + assert recorder.logged == () + assert len(recorder.failed) == 1 + assert (recorder.failed[0]["prompt_tokens"], recorder.failed[0]["completion_tokens"]) == (71, 19) + assert recorder.failed[0]["response_cost"] == pytest.approx(_round_cost(_ROUND_ONE_BODY)) + assert recorder.failed[0]["response_cost"] > 0 + assert call_tool.await_count == 1 + assert call_tool.await_args is not None + tool_logging_obj: Final = call_tool.await_args.kwargs["litellm_logging_obj"] + assert tool_logging_obj.litellm_call_id == f"{recorder.failed[0]['id']}_call_echo" + assert tool_logging_obj.litellm_trace_id == recorder.failed[0]["trace_id"] + + +_RETRIED_ROUND_ONE: Final = MappingProxyType( + { + **_ROUND_ONE_BODY, + "id": "resp_retried_round_one", + "usage": { + "input_tokens": 500, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens": 40, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 540, + }, + } +) +_SERVER_ERROR: Final = MappingProxyType( + {"error": {"message": "upstream overloaded", "type": "server_error", "code": "server_error"}} +) + + +def _ok(body: Mapping[str, object]) -> httpx.Response: + return httpx.Response(200, content=json.dumps(body, default=dict)) + + +def _server_error() -> httpx.Response: + return httpx.Response(500, content=json.dumps(_SERVER_ERROR, default=dict)) + + +def _router_retrying_once() -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": "gpt-5.6", + "litellm_params": { + "model": "openai/gpt-5.6", + "api_key": "sk-mcp-rounds", + "api_base": "https://mcp-rounds.example.com/v1", + }, + } + ], + num_retries=1, + retry_after=0, + ) + + +@pytest.mark.asyncio +async def test_non_stream_auto_execute_retried_after_a_failed_follow_up_bills_each_attempt_its_own_rounds( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_responses( + respx_mock, + monkeypatch, + (_ok(_ROUND_ONE_BODY), _server_error(), _ok(_RETRIED_ROUND_ONE), _ok(_ROUND_TWO_BODY)), + ) + recorder: Final = _RoundRecorder() + + with _echo_tool_served("PING"): + response: Final = await _router_retrying_once().aresponses( + model="gpt-5.6", + input="Call echo with PING", + tools=[dict(_ECHO_MCP_TOOL)], + litellm_logging_obj=_rounds_logging_obj(recorder, "mcp-rounds-retried"), + ) + await _settle(lambda: len(recorder.logged) > 0) + + assert isinstance(response, ResponsesAPIResponse) + assert response.usage is not None + assert len(recorder.failed) == 1 + assert (recorder.failed[0]["prompt_tokens"], recorder.failed[0]["completion_tokens"]) == (71, 19) + assert recorder.failed[0]["response_cost"] == pytest.approx(_round_cost(_ROUND_ONE_BODY)) + assert len(recorder.logged) == 1 + row: Final = recorder.logged[0] + assert row["id"] == response.id + assert (row["prompt_tokens"], row["completion_tokens"]) == (500 + 149, 40 + 5) + assert (row["prompt_tokens"], row["completion_tokens"]) == (response.usage.input_tokens, response.usage.output_tokens) + assert row["response_cost"] == pytest.approx(_round_cost(_RETRIED_ROUND_ONE) + _round_cost(_ROUND_TWO_BODY)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("retried_first_rounds", "billed_tokens"), + [ + pytest.param((), (71, 19), id="the retry fails its first call"), + pytest.param((_RETRIED_ROUND_ONE,), (71 + 500, 19 + 40), id="the retry fails its follow-up"), + ], +) +async def test_non_stream_auto_execute_failing_every_retry_bills_each_paid_first_round_once_on_the_request_failure( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + retried_first_rounds: tuple[Mapping[str, object], ...], + billed_tokens: tuple[int, int], +): + _serve_responses( + respx_mock, + monkeypatch, + (_ok(_ROUND_ONE_BODY), _server_error(), *(_ok(body) for body in retried_first_rounds), _server_error()), + ) + logging_obj: Final = _rounds_logging_obj(_RoundRecorder(), "mcp-rounds-every-retry-fails") + + with _echo_tool_served("PING"), pytest.raises(litellm.InternalServerError): + await _router_retrying_once().aresponses( + model="gpt-5.6", + input="Call echo with PING", + tools=[dict(_ECHO_MCP_TOOL)], + litellm_logging_obj=logging_obj, + ) + + failure_row_fields: Final = _failure_fields_to_lift(MappingProxyType({"litellm_logging_obj": logging_obj})) + billed_usage: Final = failure_row_fields["combined_usage_object"] + assert isinstance(billed_usage, Usage) + assert (billed_usage.prompt_tokens, billed_usage.completion_tokens) == billed_tokens + assert failure_row_fields["response_cost"] == pytest.approx( + _round_cost(_ROUND_ONE_BODY) + sum(_round_cost(body) for body in retried_first_rounds) + ) + + +@pytest.mark.asyncio +async def test_non_stream_auto_execute_tool_log_continues_the_request_trace_without_renaming_it( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_rounds(respx_mock, monkeypatch, (_ROUND_ONE_BODY, _ROUND_TWO_BODY)) + call_tool: Final = AsyncMock(return_value=CallToolResult(content=[TextContent(type="text", text="PING")])) + recorder: Final = _RoundRecorder() + + await _auto_execute_through_mcp_server(monkeypatch, call_tool, _rounds_logging_obj(recorder, "mcp-rounds-trace")) + await _settle(lambda: len(recorder.logged) > 0) + + assert call_tool.await_args is not None + tool_metadata: Final = call_tool.await_args.kwargs["litellm_logging_obj"].litellm_params["metadata"] + assert tool_metadata["existing_trace_id"] == recorder.logged[0]["trace_id"] + + +class _ListingRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.listings: tuple[Mapping[str, object], ...] = () + + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + if kwargs.get("call_type") == "list_mcp_tools": + self.listings = (*self.listings, kwargs) + + +def _serve_the_demo_mcp_listing(monkeypatch: pytest.MonkeyPatch) -> _ListingRecorder: + echo_tool: Final = MCPTool.model_validate( + {"name": "echo", "inputSchema": {"type": "object", "properties": {"value": {"type": "string"}}}}, + by_name=False, + ) + demo_server: Final = SimpleNamespace( + server_id="demo", name="demo", alias="demo", server_name="demo", auth_type=None, extra_headers=None + ) + monkeypatch.setitem( + sys.modules, "litellm.proxy.proxy_server", SimpleNamespace(proxy_logging_obj=None, prisma_client=None) + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + SimpleNamespace( + catalog=SimpleNamespace(operation=nullcontext), + get_registry=MagicMock(return_value={}), + call_tool=AsyncMock(return_value=CallToolResult(content=[TextContent(type="text", text="PING")])), + get_mcp_server_from_tool_name=MagicMock(return_value=None), + get_mcp_server_by_name=MagicMock(return_value=None), + get_allowed_mcp_servers=AsyncMock(return_value=["demo"]), + get_mcp_servers_from_ids=MagicMock(return_value=[demo_server]), + ), + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", + AsyncMock(return_value=[demo_server]), + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.catalog.get_filtered_server_tools", + AsyncMock(return_value=(ListToolsResult(tools=[echo_tool]), ServerListOk(tool_count=1))), + ) + listing_recorder: Final = _ListingRecorder() + monkeypatch.setattr(litellm, "_async_success_callback", [listing_recorder]) + return listing_recorder + + +async def _assert_the_listing_logs_in_the_request_trace( + listing_recorder: _ListingRecorder, request_recorder: _RoundRecorder +) -> None: + await _settle(lambda: len(request_recorder.logged) > 0 and len(listing_recorder.listings) > 0) + request_trace_id: Final = request_recorder.logged[0]["trace_id"] + assert len(listing_recorder.listings) == 1 + listing: Final = listing_recorder.listings[0] + listing_row: Final = cast(StandardLoggingPayload, listing["standard_logging_object"]) + assert listing_row["trace_id"] == request_trace_id + assert cast(Mapping[str, Mapping[str, object]], listing["litellm_params"])["metadata"][ + "existing_trace_id" + ] == request_trace_id + + +@pytest.mark.asyncio +async def test_non_stream_auto_execute_tool_listing_log_continues_the_request_trace( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_rounds(respx_mock, monkeypatch, (_ROUND_ONE_BODY, _ROUND_TWO_BODY)) + listing_recorder: Final = _serve_the_demo_mcp_listing(monkeypatch) + request_recorder: Final = _RoundRecorder() + + await litellm.aresponses( + model="openai/gpt-5.6", + input="Call echo with PING", + tools=[dict(_ECHO_MCP_TOOL)], + api_key="sk-mcp-rounds", + api_base="https://mcp-rounds.example.com/v1", + litellm_call_id="mcp-listing-trace", + litellm_trace_id="router-trace", + litellm_logging_obj=_rounds_logging_obj(request_recorder, "mcp-listing-trace"), + ) + + await _assert_the_listing_logs_in_the_request_trace(listing_recorder, request_recorder) + + +@pytest.mark.asyncio +async def test_bridged_auto_execute_tool_listing_log_continues_the_request_trace( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + respx_mock.post("https://mcp-rounds.example.com/v1/messages").mock( + side_effect=[ + httpx.Response( + 200, + json=_anthropic_message( + "msg_tool_round", + ({"type": "tool_use", "id": "toolu_echo", "name": "echo", "input": {"value": "PING"}},), + "tool_use", + ), + ), + httpx.Response(200, json=_anthropic_message("msg_answer", ({"type": "text", "text": "PING"},), "end_turn")), + ] + ) + listing_recorder: Final = _serve_the_demo_mcp_listing(monkeypatch) + request_recorder: Final = _RoundRecorder() + + await litellm.aresponses( + model="anthropic/claude-haiku-4-5", + input="Call echo with PING", + tools=[dict(_ECHO_MCP_TOOL)], + api_key="sk-mcp-rounds", + api_base="https://mcp-rounds.example.com", + litellm_call_id="mcp-bridged-listing-trace", + litellm_trace_id="router-trace", + litellm_logging_obj=_rounds_logging_obj(request_recorder, "mcp-bridged-listing-trace"), + ) + + await _assert_the_listing_logs_in_the_request_trace(listing_recorder, request_recorder) + + +class _CallRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logged: tuple[Mapping[str, object], ...] = () + + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self.logged = (*self.logged, kwargs) + + def of(self, call_type: str) -> tuple[Mapping[str, object], ...]: + return tuple(kwargs for kwargs in self.logged if kwargs.get("call_type") == call_type) + + +def _langfuse_trace_of(logged: Mapping[str, object]) -> object: + litellm_params: Final = cast(Mapping[str, object], logged["litellm_params"]) + return langfuse_trace_id( + metadata=langfuse_trace_metadata(litellm_params["metadata"], litellm_params.get("proxy_server_request")), + standard_trace_id=cast(StandardLoggingPayload, logged["standard_logging_object"])["trace_id"], + litellm_trace_id=logged.get("litellm_trace_id"), + litellm_call_id=logged.get("litellm_call_id"), + proxy_server_request=litellm_params.get("proxy_server_request"), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "body_metadata", + [None, {"trace_id": "client-trace", "session_id": "client-session"}], + ids=["no_body_metadata", "body_metadata_with_its_own_trace_and_session"], +) +async def test_auto_execute_mcp_logs_join_the_langfuse_trace_of_a_session_header_request( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, body_metadata: dict[str, str] | None +): + _serve_rounds(respx_mock, monkeypatch, (_ROUND_ONE_BODY, _ROUND_TWO_BODY)) + listing_recorder: Final = _serve_the_demo_mcp_listing(monkeypatch) + call_recorder: Final = _CallRecorder() + monkeypatch.setattr(litellm, "_async_success_callback", [listing_recorder, call_recorder]) + session_id: Final = "sess-abc12345" + + await litellm.aresponses( + model="openai/gpt-5.6", + input="Call echo with PING", + tools=[dict(_ECHO_MCP_TOOL)], + api_key="sk-mcp-rounds", + api_base="https://mcp-rounds.example.com/v1", + litellm_call_id="mcp-session-trace", + litellm_trace_id=session_id, + litellm_session_id=session_id, + metadata=body_metadata, + litellm_metadata={"trace_id": session_id, "session_id": session_id}, + proxy_server_request={"headers": {"x-litellm-session-id": session_id}}, + ) + await _settle(lambda: len(call_recorder.of("aresponses")) > 0 and len(listing_recorder.listings) > 0) + + parent_trace_id: Final = _langfuse_trace_of(call_recorder.of("aresponses")[0]) + call_tool: Final = cast(AsyncMock, mcp_server_manager.global_mcp_server_manager.call_tool) + assert call_tool.await_args is not None + tool_logging_obj: Final = cast(Logging, call_tool.await_args.kwargs["litellm_logging_obj"]) + mcp_log_traces: Final = ( + cast(Mapping[str, Mapping[str, object]], listing_recorder.listings[0]["litellm_params"])["metadata"][ + "existing_trace_id" + ], + cast(Mapping[str, Mapping[str, object]], tool_logging_obj.model_call_details["litellm_params"])["metadata"][ + "existing_trace_id" + ], + ) + assert parent_trace_id == "mcp-session-trace" + assert mcp_log_traces == (parent_trace_id, parent_trace_id) + + +@pytest.mark.asyncio +async def test_non_stream_auto_execute_bills_the_paid_first_round_when_the_client_disconnects_during_a_tool_call( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_rounds(respx_mock, monkeypatch, (_ROUND_ONE_BODY,)) + logging_obj: Final = _rounds_logging_obj(_RoundRecorder(), "mcp-rounds-disconnect") + + with pytest.raises(asyncio.CancelledError): + await _auto_execute_through_mcp_server( + monkeypatch, AsyncMock(side_effect=asyncio.CancelledError()), logging_obj + ) + + failure_row_fields: Final = _failure_fields_to_lift(MappingProxyType({"litellm_logging_obj": logging_obj})) + billed_usage: Final = failure_row_fields["combined_usage_object"] + assert isinstance(billed_usage, Usage) + assert (billed_usage.prompt_tokens, billed_usage.completion_tokens) == (71, 19) + assert failure_row_fields["response_cost"] == pytest.approx(_round_cost(_ROUND_ONE_BODY)) + + +@pytest.mark.asyncio +async def test_non_stream_auto_execute_bills_nothing_for_rounds_served_from_the_response_cache( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + """A cached final round hands back the id of the turn that cached it, so that turn's row is logged under a + unique id prefixed by the client's, and as a cache hit only when no round was paid.""" + monkeypatch.setattr(litellm, "cache", litellm.Cache(type="local")) + provider: Final = _serve_rounds(respx_mock, monkeypatch, (_ROUND_ONE_BODY, _ROUND_TWO_BODY, _ROUND_TWO_BODY)) + warm: Final = _RoundRecorder() + await _auto_execute(_rounds_logging_obj(warm, "mcp-cache-warm"), _ECHO_MCP_TOOL, "PING") + await _settle(lambda: len(warm.logged) > 0) + + first_cached: Final = _RoundRecorder() + first_cached_response: Final = await _auto_execute( + _rounds_logging_obj(first_cached, "mcp-cache-first-round"), + {**_ECHO_MCP_TOOL, "server_description": "echo server"}, + "PONG", + ) + all_cached: Final = _RoundRecorder() + all_cached_response: Final = await _auto_execute( + _rounds_logging_obj(all_cached, "mcp-cache-every-round"), + {**_ECHO_MCP_TOOL, "server_description": "echo server, same tools"}, + "PONG", + ) + await _settle(lambda: len(first_cached.logged) > 0 and len(all_cached.logged) > 0) + + assert provider.call_count == 3 + for logged, client in ((first_cached.logged, first_cached_response), (all_cached.logged, all_cached_response)): + assert len(logged) == 1 + assert logged[0]["id"].startswith(client.id) + assert _output_types(logged[0]["response"]["output"]) == _output_types(client.output) + assert "function_call" not in _output_types(client.output) + assert (logged[0]["prompt_tokens"], logged[0]["completion_tokens"]) == (71 + 149, 19 + 5) + assert first_cached.logged[0]["response_cost"] == pytest.approx(_round_cost(_ROUND_TWO_BODY)) + assert first_cached.logged[0]["cost_breakdown"]["total_cost"] == pytest.approx(_round_cost(_ROUND_TWO_BODY)) + assert first_cached_response.hidden_params["response_cost"] == pytest.approx(_round_cost(_ROUND_TWO_BODY)) + assert all_cached.logged[0]["response_cost"] == 0.0 + assert all_cached_response.hidden_params["response_cost"] == 0.0 + assert first_cached.logged[0]["id"] == first_cached_response.id + assert first_cached.logged[0]["cache_hit"] is None + assert all_cached.logged[0]["id"].startswith(f"{all_cached_response.id}_cache_hit") + assert all_cached.logged[0]["cache_hit"] is True + + +_PRIOR_TURN: Final = MappingProxyType( + { + "request_id": "chatcmpl-prior-turn", + "session_id": "prior-session", + "proxy_server_request": {"input": "My favorite color is teal."}, + "response": { + "id": "chatcmpl-prior-turn", + "object": "chat.completion", + "created": 1767225600, + "model": "claude-haiku-4-5", + "choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "Noted."}}], + }, + } +) + + +def _anthropic_message(message_id: str, content: Sequence[Mapping[str, object]], stop: str) -> dict[str, object]: + return { + "id": message_id, + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5", + "content": [dict(block) for block in content], + "stop_reason": stop, + "usage": {"input_tokens": 50, "output_tokens": 10}, + } + + +class _SpendLogs: + def __init__(self, rows: Sequence[Mapping[str, object]]) -> None: + self.rows: Final = tuple(rows) + self.asked: tuple[frozenset[str], ...] = () + + async def query_raw(self, query: str, request_ids: Sequence[str]) -> list[dict[str, object]]: + self.asked = (*self.asked, frozenset(request_ids)) + return [dict(row) for row in self.rows if row["request_id"] in request_ids] + + +@pytest.mark.asyncio +async def test_bridged_auto_execute_follow_up_continues_the_client_conversation_in_its_session( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + spend_logs: Final = _SpendLogs((_PRIOR_TURN,)) + monkeypatch.setitem( + sys.modules, + "litellm.proxy.proxy_server", + SimpleNamespace(proxy_logging_obj=None, prisma_client=SimpleNamespace(db=spend_logs), disable_spend_logs=False), + ) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + thinking: Final = {"type": "thinking", "thinking": "Echo PING back.", "signature": "sig-round-one"} + provider: Final = respx_mock.post("https://mcp-rounds.example.com/v1/messages").mock( + side_effect=[ + httpx.Response( + 200, + json=_anthropic_message( + "msg_tool_round", + (thinking, {"type": "tool_use", "id": "toolu_echo", "name": "echo", "input": {"value": "PING"}}), + "tool_use", + ), + ), + httpx.Response(200, json=_anthropic_message("msg_answer", ({"type": "text", "text": "PING"},), "end_turn")), + ] + ) + echo_tool: Final = MCPTool.model_validate( + {"name": "echo", "inputSchema": {"type": "object", "properties": {"value": {"type": "string"}}}}, + by_name=False, + ) + execute: Final = AsyncMock(return_value=[{"tool_call_id": "toolu_echo", "result": "PING", "name": "echo"}]) + recorder: Final = _RoundRecorder() + + with ( + patch.object( + LiteLLM_Proxy_MCP_Handler, + "process_mcp_tools_without_openai_transform", + AsyncMock(return_value=([echo_tool], {"echo": "demo"})), + ), + patch.object(LiteLLM_Proxy_MCP_Handler, "execute_tool_calls", execute), + ): + await litellm.aresponses( + model="anthropic/claude-haiku-4-5", + input="Call echo with PING", + previous_response_id="chatcmpl-prior-turn", + tools=[dict(_ECHO_MCP_TOOL)], + api_key="sk-mcp-rounds", + api_base="https://mcp-rounds.example.com", + litellm_logging_obj=_rounds_logging_obj(recorder, "mcp-rounds-bridged"), + ) + await _settle(lambda: len(recorder.logged) > 0) + + follow_up: Final = json.loads(provider.calls[1].request.content)["messages"] + assert spend_logs.asked == (frozenset({"chatcmpl-prior-turn"}),) * 2 + assert [message["role"] for message in follow_up] == ["user", "assistant", "user", "assistant", "user"] + assert json.dumps(follow_up).count("My favorite color is teal.") == 1 + assert [block["type"] for block in follow_up[3]["content"]] == ["thinking", "tool_use"] + assert follow_up[3]["content"][0]["signature"] == "sig-round-one" + assert execute.await_args is not None + assert execute.await_args.kwargs["litellm_trace_id"] == "prior-session" + assert recorder.logged[0]["trace_id"] == "prior-session" + + def _event_type(event: object) -> str: event_type: Final = getattr(event, "type", "") return str(getattr(event_type, "value", event_type)) diff --git a/tests/unit/responses/mcp/test_round_billing.py b/tests/unit/responses/mcp/test_round_billing.py new file mode 100644 index 00000000000..993b38de16a --- /dev/null +++ b/tests/unit/responses/mcp/test_round_billing.py @@ -0,0 +1,218 @@ +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +import pytest + +from litellm.responses.mcp.round_billing import billed_for_every_round, merged_round_usage, summed_cost_breakdown +from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse +from litellm.types.utils import CostBreakdown + + +def _round(response_id: str, response_cost: object, served_from_cache: bool = False) -> ResponsesAPIResponse: + response: Final = ResponsesAPIResponse.model_validate( + MappingProxyType( + { + "id": response_id, + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-5.6", + "output": (), + "usage": MappingProxyType({"input_tokens": 10, "output_tokens": 2, "total_tokens": 12}), + } + ) + ) + response.hidden_params["response_cost"] = response_cost + if served_from_cache: + response.hidden_params["cache_hit"] = True + return response + + +def test_summed_cost_breakdown_adds_each_charge_and_keeps_the_final_round_rates(): + first: Final[CostBreakdown] = { + "input_cost": 0.001, + "output_cost": 0.002, + "total_cost": 0.003, + "additional_costs": {"web_search": 0.01}, + "discount_percent": 0.1, + "margin_percent": 0.2, + "service_tier": "default", + } + final: Final[CostBreakdown] = { + "input_cost": 0.004, + "output_cost": 0.005, + "total_cost": 0.009, + "additional_costs": {"web_search": 0.01, "file_search": 0.0025}, + "discount_percent": 0.1, + "margin_percent": 0.2, + "service_tier": "default", + } + + summed: Final = summed_cost_breakdown( + first=_round("resp_first", 0.003), first_breakdown=first, final=_round("resp_final", 0.009), final_breakdown=final + ) + + assert summed is not None + assert summed["input_cost"] == pytest.approx(0.005) + assert summed["output_cost"] == pytest.approx(0.007) + assert summed["total_cost"] == pytest.approx(0.012) + assert summed["additional_costs"] == {"web_search": pytest.approx(0.02), "file_search": pytest.approx(0.0025)} + assert (summed["discount_percent"], summed["margin_percent"]) == (0.1, 0.2) + assert summed["service_tier"] == "default" + + +def test_summed_cost_breakdown_is_none_when_the_final_round_was_not_repriced(): + breakdown: Final[CostBreakdown] = {"input_cost": 0.001, "output_cost": 0.002, "total_cost": 0.003} + first: Final = _round("resp_first", 0.003) + final: Final = _round("resp_final", 0.003) + + assert summed_cost_breakdown(first=first, first_breakdown=breakdown, final=final, final_breakdown=breakdown) is None + assert summed_cost_breakdown(first=first, first_breakdown=None, final=final, final_breakdown=breakdown) is None + assert ( + summed_cost_breakdown( + first=_round("resp_first", 0.003, served_from_cache=True), + first_breakdown=breakdown, + final=final, + final_breakdown=breakdown, + ) + is None + ) + + +def test_a_final_round_served_from_the_response_cache_keeps_the_paid_first_round_breakdown(): + breakdown: Final[CostBreakdown] = {"input_cost": 0.001, "output_cost": 0.002, "total_cost": 0.003} + + assert ( + summed_cost_breakdown( + first=_round("resp_first", 0.003), + first_breakdown=breakdown, + final=_round("resp_final", 0.0, served_from_cache=True), + final_breakdown=breakdown, + ) + == breakdown + ) + + +def test_merged_round_usage_sums_token_details_and_keeps_provider_flags(): + first: Final = ResponseAPIUsage.model_validate( + MappingProxyType( + { + "input_tokens": 71, + "input_tokens_details": MappingProxyType({"cached_tokens": 32}), + "output_tokens": 19, + "total_tokens": 90, + "is_byok": True, + } + ) + ) + final: Final = ResponseAPIUsage.model_validate( + MappingProxyType( + { + "input_tokens": 149, + "input_tokens_details": MappingProxyType({"cached_tokens": 64}), + "output_tokens": 5, + "total_tokens": 154, + "is_byok": True, + } + ) + ) + + merged: Final = merged_round_usage(first, final) + + assert merged is not None + assert (merged.input_tokens, merged.output_tokens, merged.total_tokens) == (220, 24, 244) + assert merged.input_tokens_details is not None + assert merged.input_tokens_details.cached_tokens == 96 + assert merged.model_dump()["is_byok"] is True + + +def test_billed_for_every_round_leaves_the_cost_to_the_caller_when_a_round_has_no_price(): + billed: Final = billed_for_every_round( + first=_round("resp_first", None), first_breakdown=None, final=_round("resp_final", 0.004), final_breakdown=None + ) + + assert billed.id == "resp_final" + assert billed.hidden_params["response_cost"] is None + + +def test_billed_for_every_round_does_not_rebill_the_final_round_object(): + final: Final = _round("resp_final", 0.004) + + billed: Final = billed_for_every_round( + first=_round("resp_first", 0.001), first_breakdown=None, final=final, final_breakdown=None + ) + + assert billed.hidden_params["response_cost"] == pytest.approx(0.005) + assert final.hidden_params["response_cost"] == 0.004 + + +@pytest.mark.parametrize( + ("final_cached", "request_cost", "request_margin"), + [ + (False, 0.012, 0.002), + (True, 0.0032, 0.0012), + ], +) +def test_the_fixed_margin_is_charged_once_per_request_across_charged_rounds( + final_cached: bool, request_cost: float, request_margin: float +): + first_breakdown: Final[CostBreakdown] = { + "input_cost": 0.0015, + "output_cost": 0.0005, + "total_cost": 0.0032, + "margin_percent": 0.1, + "margin_fixed_amount": 0.001, + "margin_total_amount": 0.0012, + } + final_breakdown: Final[CostBreakdown] = { + "input_cost": 0.006, + "output_cost": 0.002, + "total_cost": 0.0098, + "margin_percent": 0.1, + "margin_fixed_amount": 0.001, + "margin_total_amount": 0.0018, + } + first: Final = _round("resp_first", 0.0032) + final: Final = _round("resp_final", 0.0 if final_cached else 0.0098, served_from_cache=final_cached) + + billed: Final = billed_for_every_round( + first=first, first_breakdown=first_breakdown, final=final, final_breakdown=final_breakdown + ) + breakdown: Final = summed_cost_breakdown( + first=first, first_breakdown=first_breakdown, final=final, final_breakdown=final_breakdown + ) + + assert billed.hidden_params["response_cost"] == pytest.approx(request_cost) + assert breakdown is not None + assert breakdown["total_cost"] == pytest.approx(request_cost) + assert breakdown["margin_fixed_amount"] == pytest.approx(0.001) + assert breakdown["margin_total_amount"] == pytest.approx(request_margin) + + +@pytest.mark.parametrize( + ("first_cached", "final_cached", "kept_breakdown"), + [ + (True, False, "final"), + (False, True, "first"), + (True, True, None), + ], +) +def test_a_round_served_from_the_response_cache_drops_out_of_the_breakdown( + first_cached: bool, final_cached: bool, kept_breakdown: str | None +): + breakdowns: Final[Mapping[str, CostBreakdown]] = MappingProxyType( + { + "first": {"input_cost": 0.001, "output_cost": 0.002, "total_cost": 0.003}, + "final": {"input_cost": 0.004, "output_cost": 0.005, "total_cost": 0.009}, + } + ) + + breakdown: Final = summed_cost_breakdown( + first=_round("resp_first", 0.003, served_from_cache=first_cached), + first_breakdown=breakdowns["first"], + final=_round("resp_final", 0.009, served_from_cache=final_cached), + final_breakdown=breakdowns["final"], + ) + + assert breakdown == (None if kept_breakdown is None else breakdowns[kept_breakdown]) From f8217d2a03faaa092fd0789095172f58b9e36eef Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 17:53:06 +0000 Subject: [PATCH 08/58] fix(count_tokens): apply context_management edits in the Gemini /v1/messages/count_tokens count (#46843) * fix(count_tokens): apply context_management edits in the Gemini /v1/messages/count_tokens count /v1/messages runs clear_tool_uses context edits before it sends an Anthropic body to a Gemini or Vertex AI deployment, but /v1/messages/count_tokens counted the body unedited, so a request with older tool results cleared counted about 2.8x what inference bills. The count now runs the same history and context-edit step inference runs, honors additional_drop_params on the deployment, and skips compact edits so a count never makes a paid summary call. * fix(count_tokens): match the Gemini count to inference for compact edits and drop opt-outs Run compact edits in the count without the summary call so trigger checks and history slicing match /v1/messages, honor a request-level additional_drop_params, keep the full body when any deployment in the model group drops context_management, and answer 400 naming a block Gemini cannot count when an edit trigger fails on it * fix(count_tokens): count the unedited body when an edit fails and honor deployment ids in the drop check * fix(count_tokens): log a failed edit in the count at debug, since the send already logs it * fix(count_tokens): resolve a deployment id before wildcard routes in the drop check, as the router does * fix(count_tokens): honor the router-wide additional_drop_params opt-out in the drop check The router merges default_litellm_params between the deployment and the request, so a router-wide context_management opt-out makes /v1/messages send the full history. The count now reads it in the same order. Adds the audit's wire tests for the count's context edits. * fix(count_tokens): read the opt-out of a deployment the router serves by its litellm model name * fix(count_tokens): read the opt-out and drop_params from the deployments the router and completion_model send to * fix(messages): drop an invalid context_management spec under drop_params so the send matches the count * fix(count_tokens): give Claude Code and Codex counts the drop_params default /v1/messages gives them The proxy turns on drop_params for Claude Code and Codex user agents on /v1/messages. The count now applies the same default, so an invalid context_management spec drops on both endpoints instead of the count answering 400 while the send answers 200. The wire test also covers default_fallbacks, model_group_alias, and completion_model routes. * test(count_tokens): cover a coding agent that turns drop_params off and stop two count tests running as Claude Code --- .../anthropic/count_tokens/token_counter.py | 3 +- .../pass_through/adapters/handler.py | 41 +- .../context_management/dispatcher.py | 2 + .../context_management/editors/compact.py | 3 +- .../anthropic/count_tokens/token_counter.py | 3 +- litellm/llms/base_llm/base_utils.py | 9 +- .../count_tokens/bedrock_token_counter.py | 9 +- litellm/llms/gemini/common_utils.py | 13 +- .../responses/count_tokens/token_counter.py | 3 +- litellm/llms/vertex_ai/common_utils.py | 13 +- .../llms/vertex_ai/count_tokens/context.py | 133 +++++- .../proxy/anthropic_endpoints/endpoints.py | 32 +- litellm/proxy/litellm_pre_call_utils.py | 8 +- litellm/proxy/proxy_server.py | 39 ++ litellm/router.py | 9 + litellm/types/utils.py | 7 + .../test_vertex_gemini_count_tokens_wire.py | 400 +++++++++++++++++- .../test_handler_context_management.py | 131 ++++++ .../context_management/test_compact.py | 4 +- .../vertex_ai/count_tokens/test_context.py | 184 +++++++- .../vertex_ai/test_vertex_ai_common_utils.py | 60 ++- tests/unit/proxy/test_proxy_token_counter.py | 394 ++++++++++++++++- .../test_router_get_deployments.py | 81 +++- 23 files changed, 1493 insertions(+), 88 deletions(-) create mode 100644 tests/unit/llms/anthropic/pass_through/adapters/test_handler_context_management.py diff --git a/litellm/llms/anthropic/count_tokens/token_counter.py b/litellm/llms/anthropic/count_tokens/token_counter.py index 55eaa0d2ac6..ecf050d185e 100644 --- a/litellm/llms/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/anthropic/count_tokens/token_counter.py @@ -11,7 +11,7 @@ from litellm._logging import verbose_logger from litellm.exceptions import AuthenticationError from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler from litellm.llms.base_llm.base_utils import BaseTokenCounter -from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse +from litellm.types.utils import CountContextManagement, LlmProviders, TokenCountMessagesFormat, TokenCountResponse # Global handler instance - reuse across all token counting requests anthropic_count_tokens_handler: Final = AnthropicCountTokensHandler() @@ -37,6 +37,7 @@ class AnthropicTokenCounter(BaseTokenCounter): system: JsonValue = None, messages_format: TokenCountMessagesFormat = "anthropic", output_format: JsonValue = None, + context_management: CountContextManagement | None = None, ) -> TokenCountResponse | None: """ Count tokens using Anthropic's CountTokens API. diff --git a/litellm/llms/anthropic/pass_through/adapters/handler.py b/litellm/llms/anthropic/pass_through/adapters/handler.py index 60c63964492..7b8c2550838 100644 --- a/litellm/llms/anthropic/pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/pass_through/adapters/handler.py @@ -139,7 +139,7 @@ def request_safeguards_evaluator( ) -async def _prepare_context_managed_request( +async def prepare_context_managed_request( *, model: str, messages: _AnthropicMessages, @@ -152,6 +152,8 @@ async def _prepare_context_managed_request( user_api_key_auth: "UserAPIKeyAuth | None" = None, sub_call_logging_kwargs: ProxySubCallLoggingKwargs | None = None, parent_trace_id: str | None = None, + count_only: bool = False, + drop_invalid_edits: bool = False, ) -> PolyfillResult | None: """Apply client compaction history, then optional context_management polyfill.""" from litellm.llms.anthropic.pass_through.context_management.editors.compact import ( @@ -193,6 +195,8 @@ async def _prepare_context_managed_request( user_api_key_auth=user_api_key_auth, sub_call_logging_kwargs=sub_call_logging_kwargs, parent_trace_id=parent_trace_id, + count_only=count_only, + drop_invalid_edits=drop_invalid_edits, ) if polyfill_result is not None: @@ -266,16 +270,14 @@ def _spec_has_non_compact_edits( return any(isinstance(edit.get("type"), str) and edit.get("type") != COMPACT_EDIT_TYPE for edit in edits) -def _context_management_explicitly_dropped(additional_drop_params: list[str] | None) -> bool: +def context_management_explicitly_dropped(additional_drop_params: object) -> bool: """True when the caller opted out of context_management via ``additional_drop_params``. ``drop_params`` deliberately does NOT gate the polyfill: ``context_management`` is a LiteLLM-supported param (native on Anthropic, polyfilled elsewhere), and ``drop_params`` only exists to drop genuinely unsupported params. """ - if not isinstance(additional_drop_params, list): - return False - return "context_management" in additional_drop_params + return isinstance(additional_drop_params, list) and "context_management" in additional_drop_params def _normalize_spec_edits( @@ -291,7 +293,7 @@ def _normalize_spec_edits( if not context_management_spec: return None - if _context_management_explicitly_dropped(additional_drop_params): + if context_management_explicitly_dropped(additional_drop_params): return None from litellm.llms.anthropic.pass_through.context_management.dispatcher import ( @@ -317,20 +319,23 @@ async def _run_polyfill_if_enabled( user_api_key_auth: "UserAPIKeyAuth | None" = None, sub_call_logging_kwargs: ProxySubCallLoggingKwargs | None = None, parent_trace_id: str | None = None, + count_only: bool = False, + drop_invalid_edits: bool = False, ) -> PolyfillResult | None: """Run the async context_management polyfill if a spec is present. Returns ``None`` when the spec is empty or ``context_management`` is listed in ``additional_drop_params`` (the explicit opt-out; ``drop_params`` - does not disable the polyfill because context_management is a supported + does not disable a valid spec because context_management is a supported param). Raises ``AnthropicContextManagementError`` so the /v1/messages - endpoint can emit an Anthropic-format 400. All other exceptions are + endpoint can emit an Anthropic-format 400, unless ``drop_invalid_edits`` + (``drop_params``) drops the invalid spec. All other exceptions are best-effort swallowed (matches v0 behavior). """ if not context_management_spec: return None - if _context_management_explicitly_dropped(additional_drop_params): + if context_management_explicitly_dropped(additional_drop_params): return None try: @@ -345,13 +350,17 @@ async def _run_polyfill_if_enabled( user_api_key_auth=user_api_key_auth, sub_call_logging_kwargs=sub_call_logging_kwargs, parent_trace_id=parent_trace_id, + count_only=count_only, ) except AnthropicContextManagementError: + if drop_invalid_edits: + return None # Surface validation errors so the endpoint can emit an Anthropic-format # 400. Other exception types fall into the best-effort branch below. raise except Exception as e: - verbose_logger.exception("context_management polyfill: skipping edits due to error: %s", e) + log_failure: Final = verbose_logger.debug if count_only else verbose_logger.exception + log_failure("context_management polyfill: skipping edits due to error: %s", e) # Best-effort swallow is only safe for compact-only specs, where the # caller's compaction-block-slicing safety net produces a correct # (if degraded) result. When the spec also requested non-compact @@ -709,6 +718,9 @@ class LiteLLMMessagesToCompletionTransformationHandler: """Handle non-Anthropic models asynchronously using the adapter""" context_management: Final = kwargs.pop("context_management", None) additional_drop_params: Final[list[str] | None] = kwargs.get("additional_drop_params", None) + drop_invalid_edits: Final = litellm.drop_params is True or ( + "drop_params" in kwargs and kwargs["drop_params"] is True + ) litellm_router: Final[Router | None] = _router_or_proxy_fallback(kwargs.pop("litellm_router", None)) proxy_litellm_metadata, user_api_key_auth = _extract_proxy_litellm_metadata(kwargs) @@ -716,7 +728,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: sub_call_logging_kwargs: Final = proxy_sub_call_logging_kwargs(request_kwargs) sub_call_trace_id: Final = proxy_sub_call_trace_id(request_kwargs) - polyfill_result: Final = await _prepare_context_managed_request( + polyfill_result: Final = await prepare_context_managed_request( model=model, messages=messages, tools=tools, @@ -728,6 +740,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: user_api_key_auth=user_api_key_auth, sub_call_logging_kwargs=sub_call_logging_kwargs, parent_trace_id=sub_call_trace_id, + drop_invalid_edits=drop_invalid_edits, ) effective_messages: Final = polyfill_result.messages if polyfill_result is not None else messages @@ -841,6 +854,9 @@ class LiteLLMMessagesToCompletionTransformationHandler: # bridge to it via ``run_async_function``. context_management: Final = kwargs.pop("context_management", None) additional_drop_params: Final[list[str] | None] = kwargs.get("additional_drop_params", None) + drop_invalid_edits: Final = litellm.drop_params is True or ( + "drop_params" in kwargs and kwargs["drop_params"] is True + ) # Deliberately do NOT auto-attach the proxy ``llm_router`` here: # ``run_async_function`` spawns a new event loop in a worker thread # to bridge to the async dispatcher, but the proxy router's httpx @@ -866,7 +882,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: proxy_litellm_metadata, user_api_key_auth = _extract_proxy_litellm_metadata(kwargs) request_kwargs: Final = _REQUEST_KWARGS.validate_python(kwargs) polyfill_result = run_async_function( - _prepare_context_managed_request, + prepare_context_managed_request, model=model, messages=messages, tools=tools, @@ -878,6 +894,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: user_api_key_auth=user_api_key_auth, sub_call_logging_kwargs=proxy_sub_call_logging_kwargs(request_kwargs), parent_trace_id=proxy_sub_call_trace_id(request_kwargs), + drop_invalid_edits=drop_invalid_edits, ) effective_messages: Final = polyfill_result.messages if polyfill_result is not None else messages diff --git a/litellm/llms/anthropic/pass_through/context_management/dispatcher.py b/litellm/llms/anthropic/pass_through/context_management/dispatcher.py index 27b7412292e..3e1bab09d5a 100644 --- a/litellm/llms/anthropic/pass_through/context_management/dispatcher.py +++ b/litellm/llms/anthropic/pass_through/context_management/dispatcher.py @@ -85,6 +85,7 @@ async def apply_context_management( user_api_key_auth: "UserAPIKeyAuth | None" = None, sub_call_logging_kwargs: "ProxySubCallLoggingKwargs | None" = None, parent_trace_id: str | None = None, + count_only: bool = False, ) -> PolyfillResult: """Run edits in order; return a single ``PolyfillResult``. @@ -128,6 +129,7 @@ async def apply_context_management( user_api_key_auth=user_api_key_auth, sub_call_logging_kwargs=sub_call_logging_kwargs, parent_trace_id=parent_trace_id, + count_only=count_only, ) if editor_is_async else await asyncify(editor)( diff --git a/litellm/llms/anthropic/pass_through/context_management/editors/compact.py b/litellm/llms/anthropic/pass_through/context_management/editors/compact.py index 09f49e9b663..5eb2863a241 100644 --- a/litellm/llms/anthropic/pass_through/context_management/editors/compact.py +++ b/litellm/llms/anthropic/pass_through/context_management/editors/compact.py @@ -642,6 +642,7 @@ async def apply_compact_20260112( user_api_key_auth: Optional["UserAPIKeyAuth"] = None, sub_call_logging_kwargs: ProxySubCallLoggingKwargs | None = None, parent_trace_id: str | None = None, + count_only: bool = False, ) -> PolyfillResult: """Apply ``compact_20260112``; return a ``PolyfillResult``. @@ -682,7 +683,7 @@ async def apply_compact_20260112( # Opt-in gate: no summary model configured → no-op (but still return the # Phase A-sliced/stripped messages so compaction blocks don't leak). - summary_model: Final = _read_summary_model_setting() + summary_model: Final = None if count_only else _read_summary_model_setting() if summary_model is None: applied["error"] = "summary_model_not_configured" # Slice-only forwarding: ``augmented_system`` already carries any prior diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py index 0ed432cc394..d1f45751d3e 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py @@ -13,7 +13,7 @@ from litellm.llms.azure_ai.anthropic.count_tokens.handler import ( AzureAIAnthropicCountTokensHandler, ) from litellm.llms.base_llm.base_utils import BaseTokenCounter -from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse +from litellm.types.utils import CountContextManagement, LlmProviders, TokenCountMessagesFormat, TokenCountResponse # Global handler instance - reuse across all token counting requests azure_ai_anthropic_count_tokens_handler: Final = AzureAIAnthropicCountTokensHandler() @@ -54,6 +54,7 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter): system: object | None = None, messages_format: TokenCountMessagesFormat = "anthropic", output_format: JsonValue = None, + context_management: CountContextManagement | None = None, ) -> TokenCountResponse | None: """ Count tokens using Azure AI Anthropic's CountTokens API. diff --git a/litellm/llms/base_llm/base_utils.py b/litellm/llms/base_llm/base_utils.py index 70ec6572675..5ed73e13188 100644 --- a/litellm/llms/base_llm/base_utils.py +++ b/litellm/llms/base_llm/base_utils.py @@ -16,7 +16,13 @@ from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk -from litellm.types.utils import Message, ProviderSpecificModelInfo, TokenCountMessagesFormat, TokenCountResponse +from litellm.types.utils import ( + CountContextManagement, + Message, + ProviderSpecificModelInfo, + TokenCountMessagesFormat, + TokenCountResponse, +) class BaseTokenCounter(ABC): @@ -32,6 +38,7 @@ class BaseTokenCounter(ABC): system: Any | None = None, messages_format: TokenCountMessagesFormat = "openai", output_format: JsonValue = None, + context_management: CountContextManagement | None = None, ) -> TokenCountResponse | None: pass diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index 7ff2fddbb57..97f28b9fbc5 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -16,7 +16,13 @@ from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler from litellm.llms.bedrock.count_tokens.mantle_handler import BedrockMantleCountTokensHandler from litellm.llms.bedrock_mantle.common_utils import is_mantle_claude_model from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler -from litellm.types.utils import LiteLLMPydanticObjectBase, LlmProviders, TokenCountMessagesFormat, TokenCountResponse +from litellm.types.utils import ( + CountContextManagement, + LiteLLMPydanticObjectBase, + LlmProviders, + TokenCountMessagesFormat, + TokenCountResponse, +) RUNTIME_TOKENIZER_TYPE: Final = "bedrock_api" MANTLE_TOKENIZER_TYPE: Final = "bedrock_mantle_api" @@ -109,6 +115,7 @@ class BedrockTokenCounter(BaseTokenCounter): system: object | None = None, messages_format: TokenCountMessagesFormat = "anthropic", output_format: JsonValue = None, + context_management: CountContextManagement | None = None, ) -> TokenCountResponse | None: if not messages: return None diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index 3d04954dd41..4006c89954d 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -14,7 +14,7 @@ from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues -from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse +from litellm.types.utils import CountContextManagement, TokenCountMessagesFormat, TokenCountResponse GEMINI_IMAGE_ASPECT_RATIOS: Final[dict[str, float]] = { "1:1": 1 / 1, @@ -497,12 +497,14 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): system: object | None = None, messages_format: TokenCountMessagesFormat = "openai", output_format: JsonValue = None, + context_management: CountContextManagement | None = None, ) -> TokenCountResponse | None: import copy from litellm.llms.gemini.count_tokens.handler import GoogleAIStudioTokenCounter from litellm.llms.vertex_ai.count_tokens.context import ( VertexGeminiCountContextError, + count_context_edits, count_context_error_response, gemini_count_contents, inference_request_view, @@ -511,7 +513,14 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): vertex_gemini_count_context, ) - request_view: Final = inference_request_view(messages or [], system, messages_format) + request_view: Final = await inference_request_view( + messages or [], + system, + messages_format, + count_context_edits(context_management, tools, (deployment or {}).get("litellm_params", {}), model_to_use), + ) + if isinstance(request_view, VertexGeminiCountContextError): + return count_context_error_response(request_view, request_model, model_to_use) count_context: Final = vertex_gemini_count_context(system=request_view.system, tools=tools) if isinstance(count_context, VertexGeminiCountContextError): return count_context_error_response(count_context, request_model, model_to_use) diff --git a/litellm/llms/openai/responses/count_tokens/token_counter.py b/litellm/llms/openai/responses/count_tokens/token_counter.py index 7c0697626ca..274d693baa7 100644 --- a/litellm/llms/openai/responses/count_tokens/token_counter.py +++ b/litellm/llms/openai/responses/count_tokens/token_counter.py @@ -29,7 +29,7 @@ from litellm.llms.openai.workload_identity import ( get_workload_identity_bearer_token, resolve_openai_workload_identity_config, ) -from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse +from litellm.types.utils import CountContextManagement, LlmProviders, TokenCountMessagesFormat, TokenCountResponse # Global handler instance - reuse across all token counting requests openai_count_tokens_handler: Final = OpenAICountTokensHandler() @@ -95,6 +95,7 @@ class OpenAITokenCounter(BaseTokenCounter): system: object = None, messages_format: TokenCountMessagesFormat = "openai", output_format: JsonValue = None, + context_management: CountContextManagement | None = None, ) -> TokenCountResponse | None: """ Count tokens using OpenAI's Responses API /input_tokens endpoint. diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 581e2cb15f7..27ca4acf05d 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -23,7 +23,7 @@ from litellm.types.llms.vertex_ai import ( PartType, Schema, ) -from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse +from litellm.types.utils import CountContextManagement, TokenCountMessagesFormat, TokenCountResponse from litellm.utils import supports_response_schema, supports_system_messages GEMINI_VIDEO_METADATA_KEYS: Final = MappingProxyType( @@ -1274,6 +1274,7 @@ class VertexAITokenCounter(BaseTokenCounter): system: object | None = None, messages_format: TokenCountMessagesFormat = "openai", output_format: JsonValue = None, + context_management: CountContextManagement | None = None, ) -> TokenCountResponse | None: import copy @@ -1352,6 +1353,7 @@ class VertexAITokenCounter(BaseTokenCounter): else: from litellm.llms.vertex_ai.count_tokens.context import ( VertexGeminiCountContextError, + count_context_edits, count_context_error_response, gemini_count_contents, inference_request_view, @@ -1361,7 +1363,14 @@ class VertexAITokenCounter(BaseTokenCounter): ) from litellm.llms.vertex_ai.count_tokens.handler import VertexAITokenCounter - request_view: Final = inference_request_view(messages or [], system, messages_format) + request_view: Final = await inference_request_view( + messages or [], + system, + messages_format, + count_context_edits(context_management, tools, count_tokens_params_request, model_to_use), + ) + if isinstance(request_view, VertexGeminiCountContextError): + return count_context_error_response(request_view, request_model, model_to_use) count_context: Final = vertex_gemini_count_context(system=request_view.system, tools=tools) if isinstance(count_context, VertexGeminiCountContextError): return count_context_error_response(count_context, request_model, model_to_use) diff --git a/litellm/llms/vertex_ai/count_tokens/context.py b/litellm/llms/vertex_ai/count_tokens/context.py index e094f4c5c0c..871532b7af5 100644 --- a/litellm/llms/vertex_ai/count_tokens/context.py +++ b/litellm/llms/vertex_ai/count_tokens/context.py @@ -10,11 +10,10 @@ from typing_extensions import assert_never from litellm.types.llms.anthropic import AnthropicMessagesTextParam from litellm.types.llms.openai import AllMessageValues from litellm.types.llms.vertex_ai import ContentType -from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse +from litellm.types.utils import CountContextManagement, TokenCountMessagesFormat, TokenCountResponse if TYPE_CHECKING: from litellm.exceptions import APIError - from litellm.llms.anthropic.pass_through.context_management.result import PolyfillResult @dataclass(frozen=True, slots=True) @@ -26,6 +25,7 @@ class VertexGeminiCountContext: @dataclass(frozen=True, slots=True) class VertexGeminiCountContextError: error_message: str + status_code: int = 400 _SYSTEM_BLOCKS: Final = TypeAdapter(tuple[AnthropicMessagesTextParam, ...]) @@ -39,6 +39,11 @@ _VERTEX_TOOL_KEYS: Final = MappingProxyType( ) +_INVALID_SYSTEM: Final = VertexGeminiCountContextError( + error_message="Gemini system instructions only support text blocks." +) + + def _system_parts( system: object | None, ) -> tuple[Mapping[str, JsonValue], ...] | VertexGeminiCountContextError: @@ -49,7 +54,7 @@ def _system_parts( try: blocks: Final = _SYSTEM_BLOCKS.validate_python(system) except ValidationError: - return VertexGeminiCountContextError(error_message="Gemini system instructions only support text blocks.") + return _INVALID_SYSTEM return tuple({"text": block["text"]} for block in blocks) @@ -157,7 +162,7 @@ def count_context_error_response( tokenizer_type="gemini", error=True, error_message=error.error_message, - status_code=400, + status_code=error.status_code, ) @@ -179,36 +184,120 @@ class InferenceRequestView: system: object | None -def _client_compaction(messages: Sequence[Mapping[str, object]], system: object | None) -> "PolyfillResult | None": +@dataclass(frozen=True, slots=True) +class CountContextEdits: + spec: JsonValue + model: str + tools: tuple[Mapping[str, object], ...] | None + drop_invalid_edits: bool = False + + +_CONTEXT_MANAGEMENT_SPEC: Final[TypeAdapter[Mapping[str, object] | tuple[Mapping[str, object], ...]]] = TypeAdapter( + Mapping[str, object] | tuple[Mapping[str, object], ...] +) + + +def count_context_edits( + context_management: CountContextManagement | None, + tools: Sequence[Mapping[str, object]] | None, + litellm_params: Mapping[str, object], + fallback_model: str, +) -> CountContextEdits: + model: Final = litellm_params.get("model") + return CountContextEdits( + spec=None if context_management is None else context_management.spec, + model=model if isinstance(model, str) else fallback_model, + tools=None if tools is None else tuple(tools), + drop_invalid_edits=context_management is not None and context_management.drop_invalid_edits, + ) + + +def _typed_spec(spec: JsonValue) -> Mapping[str, object] | tuple[Mapping[str, object], ...] | None: + try: + return _CONTEXT_MANAGEMENT_SPEC.validate_python(spec) + except ValidationError: + return None + + +def _typed_system( + system: object | None, +) -> str | tuple[AnthropicMessagesTextParam, ...] | None | VertexGeminiCountContextError: + if system is None or isinstance(system, str): + return system + try: + return _SYSTEM_BLOCKS.validate_python(system) + except ValidationError: + return _INVALID_SYSTEM + + +def _view_without_edits( + messages: Sequence[Mapping[str, object]], + system: object | None, + typed_system: str | Sequence[Mapping[str, object]] | None, +) -> InferenceRequestView: from litellm.llms.anthropic.pass_through.context_management.editors.compact import ( apply_client_compaction_block_history, ) - chat_messages: Final = [dict(message) for message in messages] - if system is None or isinstance(system, str): - return apply_client_compaction_block_history(messages=chat_messages, system=system) - try: - blocks: Final = _SYSTEM_BLOCKS.validate_python(system) - except ValidationError: - return None - return apply_client_compaction_block_history(messages=chat_messages, system=[dict(block) for block in blocks]) + history: Final = apply_client_compaction_block_history( + messages=[dict(message) for message in messages], + system=typed_system + if typed_system is None or isinstance(typed_system, str) + else [dict(block) for block in typed_system], + ) + if history is None: + return InferenceRequestView(messages=messages, system=system) + return InferenceRequestView(messages=history.messages, system=history.system) -def inference_request_view( - messages: Sequence[Mapping[str, object]], system: object | None, messages_format: TokenCountMessagesFormat -) -> InferenceRequestView: - """Apply the history rewrites /v1/messages runs before it converts a body for Gemini.""" +async def _anthropic_inference_view( + messages: Sequence[Mapping[str, object]], system: object | None, edits: CountContextEdits +) -> InferenceRequestView | VertexGeminiCountContextError: + from litellm.llms.anthropic.pass_through.adapters.handler import prepare_context_managed_request + from litellm.llms.anthropic.pass_through.context_management import AnthropicContextManagementError from litellm.llms.anthropic.pass_through.messages.utils import normalize_native_messages + typed_system: Final = _typed_system(system) + if isinstance(typed_system, VertexGeminiCountContextError): + return typed_system + normalized: Final = normalize_native_messages(messages) + spec: Final = _typed_spec(edits.spec) + try: + prepared: Final = await prepare_context_managed_request( + model=edits.model, + messages=normalized, + tools=None if edits.tools is None else [dict(tool) for tool in edits.tools], + system=[dict(block) for block in typed_system] if isinstance(typed_system, tuple) else typed_system, + context_management_spec=( + None if spec is None else dict(spec) if isinstance(spec, Mapping) else [dict(edit) for edit in spec] + ), + litellm_metadata=None, + additional_drop_params=None, + llm_router=None, + count_only=True, + drop_invalid_edits=edits.drop_invalid_edits, + ) + except AnthropicContextManagementError as error: + if error.status_code >= 500: + return _view_without_edits(normalized, system, typed_system) + return VertexGeminiCountContextError(error_message=error.message, status_code=error.status_code) + if prepared is None: + return InferenceRequestView(messages=normalized, system=system) + return InferenceRequestView(messages=prepared.messages, system=prepared.system) + + +async def inference_request_view( + messages: Sequence[Mapping[str, object]], + system: object | None, + messages_format: TokenCountMessagesFormat, + edits: CountContextEdits, +) -> InferenceRequestView | VertexGeminiCountContextError: + """Apply the history rewrites and context edits /v1/messages runs before it converts a body for Gemini.""" match messages_format: case "openai": return InferenceRequestView(messages=messages, system=system) case "anthropic": - normalized: Final = normalize_native_messages(messages) - compacted: Final = _client_compaction(normalized, system) - if compacted is None: - return InferenceRequestView(messages=normalized, system=system) - return InferenceRequestView(messages=compacted.messages, system=compacted.system) + return await _anthropic_inference_view(messages, system, edits) case _: assert_never(messages_format) diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index ee1ac79c820..a11268ad432 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -3,11 +3,13 @@ Unified /v1/messages endpoint - (Anthropic Spec) """ from collections.abc import Mapping +from types import MappingProxyType from typing import Final from fastapi import APIRouter, Depends, HTTPException, Request, Response from fastapi.responses import JSONResponse from pydantic import JsonValue +from typing_extensions import ReadOnly, TypedDict import litellm from litellm.anthropic_interface.exceptions import ( @@ -44,6 +46,7 @@ from litellm.proxy.common_utils.openai_error_payload import ( openai_error_type, with_litellm_call_id, ) +from litellm.proxy.litellm_pre_call_utils import should_auto_drop_params_for_agentic_cli router: Final = APIRouter() @@ -299,6 +302,26 @@ def _requested_output_format(data: Mapping[str, JsonValue]) -> JsonValue: return output_config.get("format") +_COUNTED_LITELLM_PARAMS: Final = ("additional_drop_params", "drop_params") + + +class _CountTokensBodyOptions(TypedDict): + output_format: ReadOnly[JsonValue] + context_management: ReadOnly[JsonValue] + request_litellm_params: ReadOnly[Mapping[str, JsonValue]] + + +def _count_tokens_body_options(data: Mapping[str, JsonValue], auto_drop_params: bool) -> _CountTokensBodyOptions: + sent_params: Final[Mapping[str, JsonValue]] = {key: data[key] for key in _COUNTED_LITELLM_PARAMS if key in data} + return { + "output_format": _requested_output_format(data), + "context_management": data.get("context_management"), + "request_litellm_params": MappingProxyType( + {"drop_params": True, **sent_params} if auto_drop_params else sent_params + ), + } + + @router.post( "/v1/messages/count_tokens", tags=["[beta] Anthropic Messages Token Counting"], @@ -328,7 +351,7 @@ async def count_tokens( Returns: {"input_tokens": } """ - from litellm.proxy.proxy_server import count_request_tokens + from litellm.proxy.proxy_server import count_request_tokens, proxy_config litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) try: @@ -359,7 +382,12 @@ async def count_tokens( request=token_request, call_endpoint=True, messages_format="anthropic", - output_format=_requested_output_format(data), + **_count_tokens_body_options( + data, + should_auto_drop_params_for_agentic_cli( + request.headers.get("user-agent", ""), "drop_params" in data, proxy_config + ), + ), ) except HTTPException as e: return _anthropic_error_json_response(proxy_exception_from_http_exception(e, {}), request) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 675bf0830f9..3ad5caade1b 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -945,7 +945,9 @@ def apply_missing_session_id_policy( ) -def should_auto_drop_params_for_agentic_cli(user_agent: str, data: dict, proxy_config: ProxyConfig) -> bool: +def should_auto_drop_params_for_agentic_cli( + user_agent: str, caller_set_drop_params: bool, proxy_config: ProxyConfig +) -> bool: """drop_params defaults to on for agentic CLIs so their client-specific params (e.g. Claude Code's thinking, Codex's service_tier) don't fail requests routed to providers that reject them. An explicit drop_params @@ -955,7 +957,7 @@ def should_auto_drop_params_for_agentic_cli(user_agent: str, data: dict, proxy_c if not (is_claude_code_user_agent(user_agent) or is_codex_user_agent(user_agent)): return False - if "drop_params" in data: + if caller_set_drop_params: return False config: Final = getattr(proxy_config, "config", None) litellm_settings: Final = config.get("litellm_settings") if isinstance(config, dict) else None @@ -2617,7 +2619,7 @@ async def add_litellm_data_to_request( user_agent = request.headers["user-agent"] data[_metadata_variable_name]["user_agent"] = user_agent - if should_auto_drop_params_for_agentic_cli(user_agent, data, proxy_config): + if should_auto_drop_params_for_agentic_cli(user_agent, "drop_params" in data, proxy_config): data["drop_params"] = True # Merge caller-supplied tags (x-litellm-tags header, data["tags"] root-level) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b56f0e68c4f..89fd3ae0f49 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -171,6 +171,7 @@ from litellm.router_utils.routing_groups import parse_routing_groups from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.utils import ( PRICING_OVERRIDES_KEY, + CountContextManagement, ModelResponse, ModelResponseStream, StreamingChoices, @@ -14456,6 +14457,7 @@ async def _try_provider_token_count( system: str | None = None, messages_format: TokenCountMessagesFormat = "openai", output_format: JsonValue = None, + context_management: CountContextManagement | None = None, ) -> Optional["TokenCountResponse"]: """Attempt provider-specific token counting. Returns result on success, None to fall through to local counting.""" if not provider_counter.should_use_token_counting_api(custom_llm_provider=custom_llm_provider): @@ -14471,6 +14473,7 @@ async def _try_provider_token_count( system=system, messages_format=messages_format, output_format=output_format, + context_management=context_management, ) except httpx.HTTPStatusError as e: error_message: Final = getattr(e, "message", None) or str(e) @@ -14501,6 +14504,39 @@ def _drops_params(deployment: Mapping[str, object] | None) -> bool: return litellm.drop_params is True or (litellm_params is not None and litellm_params.get("drop_params") is True) +def _router_default_litellm_params(router: Router) -> Mapping[str, object]: + defaults: Final[Mapping[str, object] | None] = getattr(router, "default_litellm_params", None) + if defaults is None: + return {} + return {key: value for key, value in defaults.items() if value is not None} + + +def _count_context_management( + context_management: JsonValue, model: str, request_litellm_params: Mapping[str, JsonValue] +) -> CountContextManagement | None: + from litellm.llms.anthropic.pass_through.adapters.handler import context_management_explicitly_dropped + + if context_management is None or llm_router is None: + return None + inference_model: Final = resolve_inference_model( + model, general_settings_view(), user_model if isinstance(user_model, str) else None + ) + router_defaults: Final = _router_default_litellm_params(llm_router) + sent_params: Final = tuple( + {**deployment["litellm_params"], **router_defaults, **request_litellm_params} + for deployment in llm_router.get_routable_deployments( + inference_model if isinstance(inference_model, str) else model + ) + ) + if any(context_management_explicitly_dropped(params.get("additional_drop_params")) for params in sent_params): + return None + return CountContextManagement( + spec=context_management, + drop_invalid_edits=litellm.drop_params is True + or any(params.get("drop_params") is True for params in sent_params), + ) + + def _system_message(system: object) -> ChatCompletionSystemMessage | None: if not isinstance(system, (str, list)) or not system: return None @@ -14531,6 +14567,8 @@ async def count_request_tokens( call_endpoint: bool = False, messages_format: TokenCountMessagesFormat = "openai", output_format: JsonValue = None, + context_management: JsonValue = None, + request_litellm_params: Mapping[str, JsonValue] = MappingProxyType({}), ) -> TokenCountResponse: prompt: Final = request.prompt messages: Final = request.messages @@ -14602,6 +14640,7 @@ async def count_request_tokens( system=system, messages_format=messages_format, output_format=output_format, + context_management=_count_context_management(context_management, request.model, request_litellm_params), ) if result is not None: return result diff --git a/litellm/router.py b/litellm/router.py index 3c969610e83..4f2358f9008 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -13727,6 +13727,15 @@ class Router: """ return [m for m in self.model_list if m["litellm_params"]["model"] == model] + def get_routable_deployments(self, model: str) -> tuple[DeploymentTypedDict, ...]: + try: + _, candidates = self._common_checks_available_deployment(model=model, request_kwargs={}) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # its signature leaves input and the result untyped + except (litellm.BadRequestError, ValueError): + return () + if isinstance(candidates, dict): + return (cast(DeploymentTypedDict, candidates),) # cast-ok: the router returns one deployment as a dict + return tuple(cast(list[DeploymentTypedDict], candidates)) # cast-ok: the router returns deployments as a list + def _try_early_resolve_deployments_for_model_not_in_names( self, model: str, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 7df8cc7452f..a97e35a9021 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2,6 +2,7 @@ import json import re import time from collections.abc import Collection, Mapping, Sequence +from dataclasses import dataclass from enum import Enum from types import MappingProxyType from typing import ( @@ -4477,6 +4478,12 @@ class LiteLLMLoggingBaseClass: TokenCountMessagesFormat: TypeAlias = Literal["openai", "anthropic"] +@dataclass(frozen=True, slots=True) +class CountContextManagement: + spec: JsonValue + drop_invalid_edits: bool + + class TokenCountResponse(LiteLLMPydanticObjectBase): total_tokens: int request_model: str diff --git a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py index 8de93256f6c..ec3fb941037 100644 --- a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py +++ b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py @@ -4,18 +4,22 @@ import threading from collections.abc import Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from contextlib import ExitStack +from dataclasses import dataclass +from itertools import chain +from pathlib import Path from types import MappingProxyType -from typing import Final +from typing import Final, Literal, TypeAlias, get_args from uuid import uuid4 import anthropic import httpx import pytest -from integration._support.client import Gateway, Scenario -from integration._support.process import graceful_stop_seconds +from integration._support.client import Gateway, Scenario, object_value, string_value +from integration._support.process import graceful_stop_seconds, owned_proxy_process from integration._support.vertex import service_account_json from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter +from typing_extensions import assert_never pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) @@ -302,6 +306,396 @@ def test_count_tokens_rejects_blocks_it_cannot_convert(gateway: Gateway, provide assert peer.drain() == () +_FILE_TEXT: Final = "line of a long file that the agent read earlier. " * 180 +_FILES: Final = 3 +_READ_FILE_TOOL: Final[JsonValue] = { + "name": "read_file", + "description": "Read a file", + "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}}, +} +_CLEAR_OLDER_TOOL_RESULTS: Final[JsonValue] = { + "edits": [ + { + "type": "clear_tool_uses_20250919", + "trigger": {"type": "input_tokens", "value": 500}, + "keep": {"type": "tool_uses", "value": 1}, + } + ] +} +_COMPACT_UNDER_THE_MINIMUM_TRIGGER: Final[JsonValue] = { + "edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 1000}}] +} + + +def _read_file_round(marker: str, index: int) -> tuple[JsonValue, ...]: + return ( + {"role": "user", "content": f"Read file {index}."}, + { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": f"toolu_{index}", "name": "read_file", "input": {"path": f"{index}.txt"}} + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": f"toolu_{index}", + "content": f"{_FILE_TEXT}", + } + ], + }, + ) + + +def _agent_body(model: str, marker: str, options: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]: + rounds: Final = tuple(chain.from_iterable(_read_file_round(marker, index) for index in range(_FILES))) + return { + "model": model, + "tools": [_READ_FILE_TOOL], + "messages": [*rounds, {"role": "user", "content": "Summarize what you read."}], + "context_management": _CLEAR_OLDER_TOOL_RESULTS, + **options, + } + + +def _files_kept(contents: JsonValue, marker: str) -> tuple[bool, ...]: + sent: Final = json.dumps(contents) + return tuple(f"" in sent for index in range(_FILES)) + + +def _contents_sent_upstream( + gateway: Gateway, peer: _GeminiPeer, body: Mapping[str, JsonValue], headers: Mapping[str, str] | None = None +) -> tuple[JsonValue, JsonValue]: + generated: Final = gateway.request("POST", "/v1/messages", {**body, "max_tokens": 16}, headers=headers) + assert generated.status_code == 200, generated.text + generate_requests: Final = tuple(request for request in peer.drain() if request.target.endswith(":generateContent")) + counted: Final = gateway.request("POST", "/v1/messages/count_tokens", body, headers=headers) + assert counted.status_code == 200, counted.text + assert _JSON_OBJECT.validate_json(counted.content)["input_tokens"] == _COUNT, counted.text + count_requests: Final = peer.drain() + assert len(generate_requests) == 1 + assert [request.target.endswith(":countTokens") for request in count_requests] == [True] + return _json_body(generate_requests[0])["contents"], _counted_contents(_json_body(count_requests[0])) + + +@pytest.mark.parametrize("provider", ("vertex_ai", "gemini")) +def test_count_tokens_counts_the_history_messages_sends_after_context_edits(gateway: Gateway, provider: str) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _provider_deployment(gateway, scenario, provider, wire.url) + marker: Final = uuid4().hex + generated, counted = _contents_sent_upstream(gateway, peer, _agent_body(model, marker, {})) + assert _files_kept(generated, marker) == (False, False, True) + assert counted == generated + + +@pytest.mark.parametrize( + ("deployment_options", "request_options", "files_kept"), + ( + ({"additional_drop_params": ["context_management"]}, {}, (True, True, True)), + ({}, {"additional_drop_params": ["context_management"]}, (True, True, True)), + ({"additional_drop_params": ["context_management"]}, {"additional_drop_params": []}, (False, False, True)), + ), + ids=("deployment-drops", "request-drops", "request-overrides-the-deployment-drop"), +) +def test_count_tokens_follows_the_context_management_opt_out_messages_follows( + gateway: Gateway, + deployment_options: Mapping[str, JsonValue], + request_options: Mapping[str, JsonValue], + files_kept: tuple[bool, ...], +) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _ai_studio_deployment(scenario, wire.url, **deployment_options) + marker: Final = uuid4().hex + generated, counted = _contents_sent_upstream(gateway, peer, _agent_body(model, marker, request_options)) + assert _files_kept(generated, marker) == files_kept + assert counted == generated + + +_RoutedCase: TypeAlias = Literal[ + "router-default-drops", + "fallback-deployment-drops", + "fallback-router-default-drops", + "fallback-request-drops", + "alias-to-a-litellm-model-drops", + "completion-model-drops", +] +_DROPS_CONTEXT_MANAGEMENT: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"additional_drop_params": ["context_management"]} +) + + +@dataclass(frozen=True, slots=True) +class _RoutedProxy: + settings: Mapping[str, JsonValue] + requested: str + request_options: Mapping[str, JsonValue] + + +def _scripted_gemini(model_name: str, model: str, api_base: str, options: Mapping[str, JsonValue]) -> JsonValue: + return { + "model_name": model_name, + "litellm_params": {"model": model, "api_base": api_base, "api_key": "scripted-gemini-key", **options}, + } + + +def _default_fallback_settings( + prefix: str, api_base: str, deployment_options: Mapping[str, JsonValue], router_defaults: Mapping[str, JsonValue] +) -> Mapping[str, JsonValue]: + return { + "model_list": [ + _scripted_gemini(f"{prefix}-fallback", f"gemini/{_MODEL}", api_base, deployment_options), + _scripted_gemini(f"{prefix}-wild/*", "gemini/*", api_base, {}), + ], + "router_settings": { + "default_fallbacks": [f"{prefix}-fallback"], + **({"default_litellm_params": dict(router_defaults)} if router_defaults else {}), + }, + } + + +def _routed_proxy(case: _RoutedCase, prefix: str, api_base: str) -> _RoutedProxy: + match case: + case "router-default-drops": + return _RoutedProxy( + { + "model_list": [_scripted_gemini(prefix, f"gemini/{_MODEL}", api_base, {})], + "router_settings": {"default_litellm_params": dict(_DROPS_CONTEXT_MANAGEMENT)}, + }, + prefix, + {}, + ) + case "fallback-deployment-drops": + return _RoutedProxy( + _default_fallback_settings(prefix, api_base, _DROPS_CONTEXT_MANAGEMENT, {}), f"{prefix}-unrouted", {} + ) + case "fallback-router-default-drops": + return _RoutedProxy( + _default_fallback_settings(prefix, api_base, {}, _DROPS_CONTEXT_MANAGEMENT), f"{prefix}-unrouted", {} + ) + case "fallback-request-drops": + return _RoutedProxy( + _default_fallback_settings(prefix, api_base, {}, {}), f"{prefix}-unrouted", _DROPS_CONTEXT_MANAGEMENT + ) + case "alias-to-a-litellm-model-drops": + return _RoutedProxy( + { + "model_list": [_scripted_gemini(prefix, f"gemini/{_MODEL}", api_base, _DROPS_CONTEXT_MANAGEMENT)], + "router_settings": {"model_group_alias": {f"{prefix}-alias": f"gemini/{_MODEL}"}}, + }, + f"{prefix}-alias", + {}, + ) + case "completion-model-drops": + return _RoutedProxy( + { + "model_list": [ + _scripted_gemini(f"{prefix}-keeps", f"gemini/{_MODEL}", api_base, {}), + _scripted_gemini(f"{prefix}-drops", f"gemini/{_MODEL}", api_base, _DROPS_CONTEXT_MANAGEMENT), + ], + "general_settings": {"completion_model": f"{prefix}-drops"}, + }, + f"{prefix}-keeps", + {}, + ) + case _: + assert_never(case) + + +@pytest.mark.parametrize("case", get_args(_RoutedCase)) +def test_count_tokens_follows_the_opt_out_of_the_deployment_proxy_settings_send_to( + gateway: Gateway, tmp_path: Path, case: _RoutedCase +) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire: + routed: Final = _routed_proxy(case, f"routed-{uuid4().hex}", wire.url) + config: Final = tmp_path / "routed.yaml" + config.write_text(json.dumps(routed.settings), encoding="utf-8") + with owned_proxy_process(gateway, tmp_path, {}, config=config) as owned: + marker: Final = uuid4().hex + generated, counted = _contents_sent_upstream( + owned.gateway, peer, _agent_body(routed.requested, marker, routed.request_options) + ) + assert _files_kept(generated, marker) == (True, True, True) + assert counted == generated + + +def _named_deployment( + scenario: Scenario, + model_name: str, + litellm_params: Mapping[str, JsonValue], + model_info: Mapping[str, JsonValue], +) -> str: + created: Final = scenario.gateway.post( + "/model/new", + {"model_name": model_name, "litellm_params": dict(litellm_params), "model_info": dict(model_info)}, + ) + identity: Final = string_value(object_value(created["model_info"])["id"]) + scenario.cleanups.callback(scenario.delete_model, identity) + return identity + + +def test_count_tokens_reads_the_opt_out_of_a_deployment_id_beside_a_wildcard_route(gateway: Gateway) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + prefix: Final = f"integration-{uuid4().hex}" + _named_deployment( + scenario, + f"{prefix}-*", + {"model": "gemini/*", "api_base": wire.url, "api_key": "scripted-gemini-key"}, + {}, + ) + deployment_id: Final = _named_deployment( + scenario, + prefix, + { + "model": f"gemini/{_MODEL}", + "api_base": wire.url, + "api_key": "scripted-gemini-key", + "additional_drop_params": ["context_management"], + }, + {"id": f"{prefix}-dropping"}, + ) + assert deployment_id == f"{prefix}-dropping" + marker: Final = uuid4().hex + generated, counted = _contents_sent_upstream(gateway, peer, _agent_body(deployment_id, marker, {})) + assert _files_kept(generated, marker) == (True, True, True) + assert counted == generated + + +@pytest.mark.parametrize( + ("requested", "wildcards", "files_kept"), + ( + ("dropping", (), (True, True, True)), + ("editing", (), (False, False, True)), + ("dropping", ("*",), (False, False, True)), + ), + ids=("deployment-drops", "sibling-deployment-drops", "wildcard-route-matches-first"), +) +def test_count_tokens_reads_the_opt_out_of_the_deployment_its_litellm_model_name_routes_to( + gateway: Gateway, requested: str, wildcards: tuple[str, ...], files_kept: tuple[bool, ...] +) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + prefix: Final = f"gemini/{_MODEL}-{uuid4().hex}" + scenario.model( + model=f"{prefix}-dropping", + api_base=wire.url, + api_key="scripted-gemini-key", + additional_drop_params=["context_management"], + ) + scenario.model(model=f"{prefix}-editing", api_base=wire.url, api_key="scripted-gemini-key") + for wildcard in wildcards: + _named_deployment( + scenario, + f"{prefix}-{wildcard}", + {"model": f"{prefix}-{wildcard}", "api_base": wire.url, "api_key": "scripted-gemini-key"}, + {}, + ) + marker: Final = uuid4().hex + generated, counted = _contents_sent_upstream(gateway, peer, _agent_body(f"{prefix}-{requested}", marker, {})) + assert _files_kept(generated, marker) == files_kept + assert counted == generated + + +def _error_type(response: httpx.Response) -> JsonValue: + return object_value(_JSON_OBJECT.validate_json(response.content)["error"])["type"] + + +def _error_message(response: httpx.Response) -> JsonValue: + return object_value(_JSON_OBJECT.validate_json(response.content)["error"])["message"] + + +@pytest.mark.parametrize("provider", ("vertex_ai", "gemini")) +def test_count_tokens_rejects_a_context_edit_messages_rejects_without_an_upstream_call( + gateway: Gateway, provider: str +) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _provider_deployment(gateway, scenario, provider, wire.url) + body: Final = _agent_body(model, uuid4().hex, {"context_management": _COMPACT_UNDER_THE_MINIMUM_TRIGGER}) + generated: Final = gateway.request("POST", "/v1/messages", {**body, "max_tokens": 16}) + counted: Final = gateway.request("POST", "/v1/messages/count_tokens", body) + assert generated.status_code == 400, generated.text + assert counted.status_code == 400, counted.text + assert _error_type(counted) == _error_type(generated) + assert peer.drain() == () + + +def test_count_tokens_rejects_an_invalid_edit_from_a_coding_agent_that_turns_drop_params_off( + gateway: Gateway, +) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _ai_studio_deployment(scenario, wire.url) + body: Final = _agent_body( + model, uuid4().hex, {"context_management": _COMPACT_UNDER_THE_MINIMUM_TRIGGER, "drop_params": False} + ) + claude_code: Final = MappingProxyType({"User-Agent": "claude-cli/2.1.0 (external, cli)"}) + generated: Final = gateway.request("POST", "/v1/messages", {**body, "max_tokens": 16}, headers=claude_code) + counted: Final = gateway.request("POST", "/v1/messages/count_tokens", body, headers=claude_code) + assert generated.status_code == 400, generated.text + assert counted.status_code == 400, counted.text + assert _error_type(counted) == _error_type(generated) + assert peer.drain() == () + + +@pytest.mark.parametrize( + ("deployment_options", "request_options", "user_agent"), + ( + ({"drop_params": True}, {}, "python-httpx"), + ({}, {"drop_params": True}, "python-httpx"), + ({}, {}, "claude-cli/2.1.0 (external, cli)"), + ({}, {}, "codex_cli_rs/0.50.0"), + ), + ids=("deployment-drop-params", "request-drop-params", "claude-code-client", "codex-client"), +) +def test_count_tokens_counts_the_unedited_body_messages_sends_when_drop_params_drops_an_invalid_edit( + gateway: Gateway, + deployment_options: Mapping[str, JsonValue], + request_options: Mapping[str, JsonValue], + user_agent: str, +) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _ai_studio_deployment(scenario, wire.url, **deployment_options) + marker: Final = uuid4().hex + generated, counted = _contents_sent_upstream( + gateway, + peer, + _agent_body(model, marker, {"context_management": _COMPACT_UNDER_THE_MINIMUM_TRIGGER, **request_options}), + {"User-Agent": user_agent}, + ) + assert _files_kept(generated, marker) == (True, True, True) + assert counted == generated + + +def test_count_tokens_answers_for_the_unedited_body_when_a_context_edit_fails_inside_the_proxy( + gateway: Gateway, +) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _ai_studio_deployment(scenario, wire.url) + unedited: Final = gateway.request( + "POST", "/v1/messages/count_tokens", {"model": model, "messages": _UNCONVERTIBLE_MESSAGES} + ) + edited: Final = gateway.request( + "POST", + "/v1/messages/count_tokens", + {"model": model, "messages": _UNCONVERTIBLE_MESSAGES, "context_management": _CLEAR_OLDER_TOOL_RESULTS}, + ) + assert unedited.status_code == 400, unedited.text + assert (edited.status_code, _error_type(edited), _error_message(edited)) == ( + unedited.status_code, + _error_type(unedited), + _error_message(unedited), + ) + assert peer.drain() == () + + def test_count_tokens_keeps_each_text_block_in_system_instruction(gateway: Gateway) -> None: peer: Final = _GeminiPeer() with wire_server(peer) as wire, gateway.scenario() as scenario: diff --git a/tests/unit/llms/anthropic/pass_through/adapters/test_handler_context_management.py b/tests/unit/llms/anthropic/pass_through/adapters/test_handler_context_management.py new file mode 100644 index 00000000000..f906530c8d9 --- /dev/null +++ b/tests/unit/llms/anthropic/pass_through/adapters/test_handler_context_management.py @@ -0,0 +1,131 @@ +from collections.abc import Mapping +from typing import Final +from unittest.mock import patch + +import pytest + +import litellm +from litellm.llms.anthropic.pass_through.adapters.handler import LiteLLMMessagesToCompletionTransformationHandler +from litellm.llms.anthropic.pass_through.context_management import AnthropicContextManagementError +from litellm.router import Router +from litellm.types.utils import Choices, Message, ModelResponse, Usage + +_TOOL_OUTPUT: Final = "line of the file under review " * 40 +_INVALID_COMPACT_EDIT: Final = { + "edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 1000}}] +} + + +def _response() -> ModelResponse: + return ModelResponse( + choices=[Choices(index=0, finish_reason="stop", message=Message(content="done"))], + usage=Usage(prompt_tokens=10, completion_tokens=1, total_tokens=11), + ) + + +def _messages() -> list[dict[str, object]]: + return [ + {"role": "user", "content": "review a.py"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "toolu_1", "name": "read_file", "input": {"path": "a.py"}}], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_1", "content": _TOOL_OUTPUT}]}, + ] + + +class _Backend: + def __init__(self) -> None: + self.calls: list[Mapping[str, object]] = [] + + async def __call__(self, **kwargs: object) -> ModelResponse: + self.calls.append(kwargs) + return _response() + + +class _SyncBackend: + def __init__(self) -> None: + self.calls: list[Mapping[str, object]] = [] + + def __call__(self, **kwargs: object) -> ModelResponse: + self.calls.append(kwargs) + return _response() + + +def _router(deployment_params: Mapping[str, object], router_defaults: Mapping[str, object]) -> Router: + return Router( + model_list=[ + { + "model_name": "gemini-agent", + "litellm_params": {"model": "gemini/gemini-2.5-flash", "api_key": "test-key", **deployment_params}, + } + ], + default_litellm_params=dict(router_defaults), + ) + + +async def _send(router: Router, request_options: Mapping[str, object], backend: _Backend) -> object: + with patch("litellm.acompletion", new=backend): + return await router.aanthropic_messages( + model="gemini-agent", + max_tokens=16, + messages=_messages(), + tools=[{"name": "read_file", "input_schema": {"type": "object", "properties": {}}}], + context_management=_INVALID_COMPACT_EDIT, + **request_options, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("deployment_params", "router_defaults", "request_options", "litellm_drop_params"), + ( + ({"drop_params": True}, {}, {}, False), + ({}, {"drop_params": True}, {}, False), + ({}, {}, {"drop_params": True}, False), + ({}, {}, {}, True), + ), + ids=("deployment-drop-params", "router-default-drop-params", "request-drop-params", "litellm-drop-params"), +) +async def test_messages_sends_the_unedited_body_when_drop_params_drops_an_invalid_edit( + deployment_params: Mapping[str, object], + router_defaults: Mapping[str, object], + request_options: Mapping[str, object], + litellm_drop_params: bool, +) -> None: + backend: Final = _Backend() + with patch.object(litellm, "drop_params", litellm_drop_params): + response: Final = await _send(_router(deployment_params, router_defaults), request_options, backend) + + assert isinstance(response, dict) + assert response["content"] == [{"type": "text", "text": "done"}] + assert len(backend.calls) == 1 + assert _TOOL_OUTPUT in str(backend.calls[0]["messages"]) + + +@pytest.mark.asyncio +async def test_messages_rejects_an_invalid_edit_without_drop_params() -> None: + backend: Final = _Backend() + with patch.object(litellm, "drop_params", False), pytest.raises(AnthropicContextManagementError) as error: + await _send(_router({}, {}), {}, backend) + + assert error.value.status_code == 400 + assert backend.calls == [] + + +def test_sync_messages_sends_the_unedited_body_when_drop_params_drops_an_invalid_edit() -> None: + backend: Final = _SyncBackend() + with patch.object(litellm, "drop_params", False), patch("litellm.completion", new=backend): + response: Final = LiteLLMMessagesToCompletionTransformationHandler.anthropic_messages_handler( + max_tokens=16, + messages=_messages(), + model="gemini/gemini-2.5-flash", + api_key="test-key", + context_management=_INVALID_COMPACT_EDIT, + drop_params=True, + ) + + assert isinstance(response, dict) + assert response["content"] == [{"type": "text", "text": "done"}] + assert len(backend.calls) == 1 + assert _TOOL_OUTPUT in str(backend.calls[0]["messages"]) diff --git a/tests/unit/llms/anthropic/pass_through/context_management/test_compact.py b/tests/unit/llms/anthropic/pass_through/context_management/test_compact.py index 0a5fbcdf63e..5e55c0ade32 100644 --- a/tests/unit/llms/anthropic/pass_through/context_management/test_compact.py +++ b/tests/unit/llms/anthropic/pass_through/context_management/test_compact.py @@ -2544,7 +2544,7 @@ async def test_prepare_context_managed_request_forwards_proxy_litellm_metadata() the summary subcall lands on the router with no parent attribution, and those tokens go unbilled to the caller's key/team.""" from litellm.llms.anthropic.pass_through.adapters.handler import ( - _prepare_context_managed_request, + prepare_context_managed_request, ) captured_summary_metadata: Dict[str, Any] = {} @@ -2561,7 +2561,7 @@ async def test_prepare_context_managed_request_forwards_proxy_litellm_metadata() ), patch("litellm.token_counter", return_value=200_000), ): - result = await _prepare_context_managed_request( + result = await prepare_context_managed_request( model=MODEL, messages=_simple_messages(), tools=None, diff --git a/tests/unit/llms/vertex_ai/count_tokens/test_context.py b/tests/unit/llms/vertex_ai/count_tokens/test_context.py index 06dacedde67..405b7234c81 100644 --- a/tests/unit/llms/vertex_ai/count_tokens/test_context.py +++ b/tests/unit/llms/vertex_ai/count_tokens/test_context.py @@ -1,5 +1,6 @@ import asyncio import base64 +import logging import time import uuid from collections.abc import Generator @@ -18,10 +19,13 @@ from litellm.llms.anthropic.pass_through.adapters.transformation import ( truncate_tool_name, ) from litellm.llms.vertex_ai.count_tokens.context import ( + CountContextEdits, + InferenceRequestView, VertexGeminiCountContext, VertexGeminiCountContextError, gemini_count_contents, inference_request_view, + unconvertible_messages_error, vertex_gemini_count_context, ) from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig @@ -31,6 +35,7 @@ _SCHEMA: Final[JsonValue] = { "properties": {}, "additionalProperties": False, } +_NO_EDITS: Final = CountContextEdits(spec=None, model="gemini/gemini-2.5-flash", tools=None) _FUNCTION_DECLARATION: Final = { "functionDeclarations": [ { @@ -383,8 +388,8 @@ async def test_image_that_cannot_be_fetched_returns_an_error_value( @pytest.mark.parametrize( "system", ("Be brief.", [{"type": "text", "text": "Be brief."}]), ids=("string-system", "text-block-system") ) -def test_anthropic_view_keeps_only_the_turns_after_a_client_compaction_block(system: JsonValue) -> None: - view: Final = inference_request_view( +async def test_anthropic_view_keeps_only_the_turns_after_a_client_compaction_block(system: JsonValue) -> None: + view: Final = await inference_request_view( ( {"role": "user", "content": "first question " * 50}, {"role": "assistant", "content": [{"type": "compaction", "content": "Summary of the earlier turns."}]}, @@ -392,16 +397,18 @@ def test_anthropic_view_keeps_only_the_turns_after_a_client_compaction_block(sys ), system, "anthropic", + _NO_EDITS, ) + assert isinstance(view, InferenceRequestView) assert view.messages == [{"role": "user", "content": "latest question"}] assert "Summary of the earlier turns." in str(view.system) assert "Be brief." in str(view.system) -def test_anthropic_view_flattens_a_replayed_web_search_result_into_text() -> None: +async def test_anthropic_view_flattens_a_replayed_web_search_result_into_text() -> None: snippet: Final = "LiteLLM ships a weekly release candidate and a stable line." - view: Final = inference_request_view( + view: Final = await inference_request_view( ( {"role": "user", "content": "What is new?"}, { @@ -426,18 +433,21 @@ def test_anthropic_view_flattens_a_replayed_web_search_result_into_text() -> Non ), None, "anthropic", + _NO_EDITS, ) + assert isinstance(view, InferenceRequestView) assistant_blocks: Final = TypeAdapter(list[dict[str, str]]).validate_python(view.messages[1]["content"]) assert [block["type"] for block in assistant_blocks] == ["text"] assert snippet in assistant_blocks[0]["text"] -def test_openai_view_is_the_request_unchanged() -> None: +async def test_openai_view_is_the_request_unchanged() -> None: messages: Final = ({"role": "user", "content": [{"type": "compaction", "content": "kept"}]},) - view: Final = inference_request_view(messages, "Be brief.", "openai") + view: Final = await inference_request_view(messages, "Be brief.", "openai", _NO_EDITS) + assert isinstance(view, InferenceRequestView) assert view.messages == messages assert view.system == "Be brief." @@ -470,7 +480,7 @@ async def test_gcs_media_type_lookup_uses_the_deployment_vertex_project() -> Non @pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) -def test_anthropic_replayed_web_search_result_is_counted_as_its_text( +async def test_anthropic_replayed_web_search_result_is_counted_as_its_text( custom_llm_provider: Literal["vertex_ai", "gemini"], ) -> None: snippet: Final = "LiteLLM ships a weekly release candidate and a stable line." @@ -498,9 +508,163 @@ def test_anthropic_replayed_web_search_result_is_counted_as_its_text( {"role": "user", "content": "Thanks."}, ) - view: Final = inference_request_view(messages, None, "anthropic") - contents: Final = asyncio.run( - gemini_count_contents(view.messages, "anthropic", "gemini-2.5-flash", custom_llm_provider, {}) + view: Final = await inference_request_view(messages, None, "anthropic", _NO_EDITS) + assert isinstance(view, InferenceRequestView) + contents: Final = await gemini_count_contents( + view.messages, "anthropic", "gemini-2.5-flash", custom_llm_provider, {} ) assert snippet in str(contents) + + +_TOOL_OUTPUT: Final = "line of a long file that the agent read earlier. " * 180 +_READ_FILE_TOOL: Final[dict[str, object]] = { + "name": "read_file", + "description": "Read a file", + "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}}, +} +_CLEAR_EDIT: Final[JsonValue] = { + "type": "clear_tool_uses_20250919", + "trigger": {"type": "input_tokens", "value": 500}, + "keep": {"type": "tool_uses", "value": 1}, +} +_CLEAR_OLD_TOOL_RESULTS: Final[JsonValue] = {"edits": [_CLEAR_EDIT]} + + +def _agent_turns(rounds: int, first_round: int = 0) -> tuple[dict[str, JsonValue], ...]: + return ( + *( + turn + for index in range(first_round, first_round + rounds) + for turn in ( + {"role": "user", "content": f"Read file {index}."}, + { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": f"toolu_{index}", "name": "read_file", "input": {"path": f"{index}"}} + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": f"toolu_{index}", "content": _TOOL_OUTPUT}], + }, + ) + ), # comprehension-ok: flattens each round into its three turns + {"role": "user", "content": "Summarize what you read."}, + ) + + +def _tool_outputs_kept(view: InferenceRequestView) -> tuple[bool, ...]: + return tuple( + _TOOL_OUTPUT in str(message["content"]) for message in view.messages if "tool_result" in str(message["content"]) + ) + + +def _edits(spec: JsonValue) -> CountContextEdits: + return CountContextEdits(spec=spec, model="gemini/gemini-2.5-flash", tools=(_READ_FILE_TOOL,)) + + +async def test_anthropic_view_clears_the_tool_results_inference_clears() -> None: + view: Final = await inference_request_view(_agent_turns(3), None, "anthropic", _edits(_CLEAR_OLD_TOOL_RESULTS)) + + assert isinstance(view, InferenceRequestView) + assert _tool_outputs_kept(view) == (False, False, True) + + +async def test_anthropic_view_applies_clear_edits_without_calling_the_compaction_summary_model() -> None: + spec: Final[JsonValue] = { + "edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 50000}}, _CLEAR_EDIT] + } + turns: Final = ( + {"role": "user", "content": "first question " * 30000}, + {"role": "assistant", "content": "first answer"}, + *_agent_turns(3), + ) + + with ( + patch( + "litellm.llms.anthropic.pass_through.context_management.editors.compact._read_summary_model_setting", + return_value="gemini/gemini-2.5-flash", + ), + patch.object(litellm, "acompletion", side_effect=AssertionError("count_tokens must not summarize")) as summary, + ): + view: Final = await inference_request_view(turns, None, "anthropic", _edits(spec)) + + assert summary.call_count == 0 + assert isinstance(view, InferenceRequestView) + assert _tool_outputs_kept(view) == (False, False, True) + + +async def test_anthropic_view_counts_the_unedited_body_without_an_error_log_when_an_edit_fails( + caplog: pytest.LogCaptureFixture, +) -> None: + with ( + patch.object(litellm, "token_counter", side_effect=RuntimeError("tokenizer unavailable")), + caplog.at_level(logging.DEBUG, logger="LiteLLM"), + ): + view: Final = await inference_request_view(_agent_turns(3), None, "anthropic", _edits(_CLEAR_OLD_TOOL_RESULTS)) + + assert isinstance(view, InferenceRequestView) + assert _tool_outputs_kept(view) == (True, True, True) + assert any("skipping edits" in record.getMessage() for record in caplog.records) + assert not any(record.levelno >= logging.ERROR for record in caplog.records) + + +async def test_anthropic_view_counts_the_turns_after_a_compaction_block_when_an_edit_before_compact_fails() -> None: + clear_before_compact: Final[JsonValue] = { + "edits": [ + { + "type": "clear_tool_uses_20250919", + "trigger": {"type": "tool_uses", "value": 3}, + "keep": {"type": "tool_uses", "value": 1}, + }, + {"type": "compact_20260112"}, + ] + } + turns: Final = ( + *_agent_turns(3)[:-1], + {"role": "assistant", "content": [{"type": "compaction", "content": "Read files 0 to 2."}]}, + *_agent_turns(2, first_round=3), + ) + + view: Final = await inference_request_view(turns, None, "anthropic", _edits(clear_before_compact)) + + assert isinstance(view, InferenceRequestView) + assert _tool_outputs_kept(view) == (True, True) + assert "Read files 0 to 2." in str(view.system) + + +async def test_anthropic_view_rejects_a_compact_trigger_inference_rejects() -> None: + low_trigger: Final[JsonValue] = { + "edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 1000}}] + } + + view: Final = await inference_request_view(_agent_turns(1), None, "anthropic", _edits(low_trigger)) + + assert isinstance(view, VertexGeminiCountContextError) + assert view.status_code == 400 + + +async def test_anthropic_view_names_the_block_gemini_cannot_count_when_an_edit_trigger_cannot_tokenize_it() -> None: + turns: Final = ( + *_agent_turns(1)[:-1], + { + "role": "user", + "content": [ + { + "type": "search_result", + "source": "https://docs.litellm.ai", + "title": "Docs", + "content": [{"type": "text", "text": "LiteLLM docs"}], + } + ], + }, + ) + + view: Final = await inference_request_view(turns, None, "anthropic", _edits(_CLEAR_OLD_TOOL_RESULTS)) + + assert isinstance(view, InferenceRequestView) + blocks_error: Final = unconvertible_messages_error(view.messages, "anthropic") + assert blocks_error is not None + assert blocks_error.status_code == 400 + assert "search_result" in blocks_error.error_message diff --git a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py index d244fe6bd90..97c0a657ca4 100644 --- a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -28,7 +28,7 @@ from litellm.llms.vertex_ai.common_utils import ( vertex_request_labels_from_litellm_params, ) from litellm.llms.vertex_ai.count_tokens import handler as vertex_count_tokens_handler -from litellm.types.utils import TokenCountMessagesFormat +from litellm.types.utils import CountContextManagement, TokenCountMessagesFormat @pytest.mark.asyncio @@ -62,7 +62,9 @@ class _VertexCountTokensPeer: return httpx.Response(200, json={"totalTokens": 37}, request=request) return httpx.Response( self.status_code, - json={"error": {"code": self.status_code, "message": "Invalid function name", "status": "INVALID_ARGUMENT"}}, + json={ + "error": {"code": self.status_code, "message": "Invalid function name", "status": "INVALID_ARGUMENT"} + }, request=request, ) @@ -133,6 +135,60 @@ async def test_vertex_gemini_count_rejects_an_image_url_it_refuses_to_fetch_with assert peer.requests == () +_TOOL_OUTPUT: Final = "line of a long file that the agent read earlier. " * 180 + + +def _tool_round(index: int) -> tuple[dict[str, object], ...]: + return ( + {"role": "user", "content": f"Read file {index}."}, + { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": f"toolu_{index}", "name": "read_file", "input": {"path": f"{index}"}} + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": f"toolu_{index}", "content": _TOOL_OUTPUT}], + }, + ) + + +@pytest.mark.asyncio +async def test_vertex_gemini_count_applies_the_context_edits_inference_applies() -> None: + with _vertex_count_tokens_peer(200) as peer: + result: Final = await VertexAITokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=[*_tool_round(0), *_tool_round(1), *_tool_round(2), {"role": "user", "content": "Summarize."}], + contents=None, + deployment={ + "litellm_params": {**_VERTEX_DEPLOYMENT["litellm_params"], "model": "vertex_ai/gemini-2.5-flash"} + }, + request_model="vertex_ai/gemini-2.5-flash", + tools=[ + {"name": "read_file", "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}}} + ], + messages_format="anthropic", + context_management=CountContextManagement( + spec={ + "edits": [ + { + "type": "clear_tool_uses_20250919", + "trigger": {"type": "input_tokens", "value": 500}, + "keep": {"type": "tool_uses", "value": 1}, + } + ] + }, + drop_invalid_edits=False, + ), + ) + + assert result is not None + assert result.total_tokens == 37 + assert len(peer.requests) == 1 + assert peer.requests[0].content.decode().count(_TOOL_OUTPUT) == 1 + + @pytest.mark.parametrize("location", ["us", "eu", "us-central1", "europe-west1", "global"]) def test_validate_vertex_location_accepts_valid(location): assert validate_vertex_location(location) == location diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 7d70e058396..f7d1970fdae 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -3,6 +3,7 @@ import json +from collections.abc import Callable, Mapping from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -153,6 +154,7 @@ async def test_anthropic_messages_count_tokens_endpoint(): # Mock request object mock_request = MagicMock(spec=Request) + mock_request.headers = {} mock_request_data = { "model": "claude-3-sonnet-20240229", "messages": [{"role": "user", "content": "Hello Claude!"}], @@ -172,7 +174,14 @@ async def test_anthropic_messages_count_tokens_endpoint(): anthropic_endpoints.read_request_body = mock_read_request_body # Mock the internal token_counter function to return a controlled response - async def mock_token_counter(request, call_endpoint=False, messages_format="openai", output_format=None): + async def mock_token_counter( + request, + call_endpoint=False, + messages_format="openai", + output_format=None, + context_management=None, + request_litellm_params=None, + ): assert call_endpoint == True, "Should be called with call_endpoint=True for Anthropic endpoint" assert request.model == "claude-3-sonnet-20240229" assert request.messages == [{"role": "user", "content": "Hello Claude!"}] @@ -225,6 +234,7 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model(): # Mock request object mock_request = MagicMock(spec=Request) + mock_request.headers = {} mock_request_data = { "model": "gpt-4", "messages": [{"role": "user", "content": "Hello GPT!"}], @@ -244,7 +254,14 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model(): anthropic_endpoints.read_request_body = mock_read_request_body # Mock the internal token_counter function to return a controlled response - async def mock_token_counter(request, call_endpoint=True, messages_format="openai", output_format=None): + async def mock_token_counter( + request, + call_endpoint=True, + messages_format="openai", + output_format=None, + context_management=None, + request_litellm_params=None, + ): assert call_endpoint == True, "Should be called with call_endpoint=True for Anthropic endpoint" assert request.model == "gpt-4" assert request.messages == [{"role": "user", "content": "Hello GPT!"}] @@ -1055,7 +1072,14 @@ async def test_anthropic_endpoint_returns_anthropic_error_format(): original_token_counter = proxy_server.count_request_tokens # Mock token_counter to raise ProxyException with Bedrock-style error - async def mock_token_counter_error(request, call_endpoint=False, messages_format="openai", output_format=None): + async def mock_token_counter_error( + request, + call_endpoint=False, + messages_format="openai", + output_format=None, + context_management=None, + request_litellm_params=None, + ): raise ProxyException( message='{"detail":{"message":"Input is too long for requested model."}}', type="token_counting_error", @@ -1105,7 +1129,14 @@ async def test_anthropic_endpoint_403_permission_error_format(): original_token_counter = proxy_server.count_request_tokens # Mock token_counter to raise ProxyException with 403 error - async def mock_token_counter_error(request, call_endpoint=False, messages_format="openai", output_format=None): + async def mock_token_counter_error( + request, + call_endpoint=False, + messages_format="openai", + output_format=None, + context_management=None, + request_litellm_params=None, + ): raise ProxyException( message='{"Message":"Bearer Token has expired"}', type="token_counting_error", @@ -1155,7 +1186,14 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): original_token_counter = proxy_server.count_request_tokens # Mock token_counter to raise ProxyException with 429 error - async def mock_token_counter_error(request, call_endpoint=False, messages_format="openai", output_format=None): + async def mock_token_counter_error( + request, + call_endpoint=False, + messages_format="openai", + output_format=None, + context_management=None, + request_litellm_params=None, + ): raise ProxyException( message="Rate limit exceeded", type="token_counting_error", @@ -1179,6 +1217,310 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): proxy_server.count_request_tokens = original_token_counter +_GEMINI_TOOL_OUTPUT: Final = "line of a long file that the agent read earlier. " * 180 + + +class _GeminiCountTokensPeer: + def __init__(self) -> None: + self.request_bodies: tuple[bytes, ...] = () + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.request_bodies = (*self.request_bodies, request.content) + return httpx.Response(200, json={"totalTokens": 41}, request=request) + + +def _agent_turns_with_tool_results(rounds: int) -> list[dict[str, JsonValue]]: + return [ + *( + turn + for index in range(rounds) + for turn in ( + {"role": "user", "content": f"Read file {index}."}, + { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": f"toolu_{index}", "name": "read_file", "input": {"path": f"{index}"}} + ], + }, + { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": f"toolu_{index}", "content": _GEMINI_TOOL_OUTPUT} + ], + }, + ) + ), # comprehension-ok: flattens each round into its three turns + {"role": "user", "content": "Summarize what you read."}, + ] + + +_READ_FILE_TOOL: Final[dict[str, JsonValue]] = { + "name": "read_file", + "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}}, +} +_CLEAR_TOOL_USES: Final[dict[str, JsonValue]] = { + "edits": [ + { + "type": "clear_tool_uses_20250919", + "trigger": {"type": "input_tokens", "value": 500}, + "keep": {"type": "tool_uses", "value": 1}, + } + ] +} +_INVALID_COMPACT_EDIT: Final[dict[str, JsonValue]] = { + "edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 1000}}] +} + + +def _gemini_deployment(model_name: str, litellm_params: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return { + "model_name": model_name, + "litellm_params": {"model": "gemini/gemini-2.5-flash", "api_key": "test-key", **litellm_params}, + } + + +async def _count_through_the_messages_endpoint( + router: Router, request_options: Mapping[str, JsonValue] +) -> tuple[object, tuple[bytes, ...]]: + import litellm.proxy.anthropic_endpoints.endpoints as anthropic_endpoints + import litellm.proxy.proxy_server as proxy_server + from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler + + peer: Final = _GeminiCountTokensPeer() + request_body: Final[dict[str, JsonValue]] = { + "messages": _agent_turns_with_tool_results(3), + "tools": [_READ_FILE_TOOL], + "context_management": _CLEAR_TOOL_USES, + **request_options, + } + mock_request: Final = MagicMock(spec=Request) + mock_request.headers = {} + + with ( + patch.object(proxy_server, "llm_router", router), + patch.object(anthropic_endpoints, "read_request_body", AsyncMock(return_value=request_body)), + patch.object( + gemini_count_tokens_handler, + "get_async_httpx_client", + return_value=AsyncHTTPHandler(transport=httpx.MockTransport(peer)), + ), + ): + response: Final = await anthropic_count_tokens(mock_request, MagicMock()) + return response, peer.request_bodies + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ( + "requested_model", + "additional_drop_params", + "router_additional_drop_params", + "request_options", + "expected_tool_outputs_sent", + ), + ( + ("gemini-agent", None, None, {}, 1), + ("gemini-agent", ["context_management"], None, {}, 3), + ("gemini-agent", None, None, {"additional_drop_params": ["context_management"]}, 3), + ("gemini-agent", ["context_management"], None, {"additional_drop_params": []}, 1), + ("gemini-agent", None, ["context_management"], {}, 3), + ("gemini-agent", None, ["context_management"], {"additional_drop_params": []}, 1), + ("gemini/gemini-2.5-flash", None, None, {}, 1), + ("gemini/gemini-2.5-flash", ["context_management"], None, {}, 3), + ("gemini/gemini-2.5-flash", None, ["context_management"], {}, 3), + ("gemini/gemini-2.5-flash", None, None, {"additional_drop_params": ["context_management"]}, 3), + ("gemini/gemini-2.5-flash", ["context_management"], None, {"additional_drop_params": []}, 1), + ), + ids=( + "edits-applied", + "deployment-drops-context-management", + "request-drops-context-management", + "request-overrides-the-deployment-drop", + "router-default-drops-context-management", + "request-overrides-the-router-default-drop", + "litellm-model-name-edits-applied", + "litellm-model-name-deployment-drops-context-management", + "litellm-model-name-router-default-drops-context-management", + "litellm-model-name-request-drops-context-management", + "litellm-model-name-request-overrides-the-deployment-drop", + ), +) +async def test_gemini_messages_count_tokens_applies_the_context_edits_inference_applies( + requested_model: str, + additional_drop_params: list[str] | None, + router_additional_drop_params: list[str] | None, + request_options: dict[str, JsonValue], + expected_tool_outputs_sent: int, +) -> None: + router: Final = Router( + model_list=[_gemini_deployment("gemini-agent", {"additional_drop_params": additional_drop_params})], + default_litellm_params={"additional_drop_params": router_additional_drop_params}, + ) + response, sent_bodies = await _count_through_the_messages_endpoint( + router, {**request_options, "model": requested_model} + ) + + assert response == {"input_tokens": 41} + assert len(sent_bodies) == 1 + assert sent_bodies[0].decode().count(_GEMINI_TOOL_OUTPUT) == expected_tool_outputs_sent + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("requested_model", "model_list"), + ( + ( + "gemini-agent", + [ + { + "model_name": "gemini-agent", + "litellm_params": {"model": "gemini/gemini-2.5-flash", "api_key": "test-key", "order": 1}, + }, + { + "model_name": "gemini-agent", + "litellm_params": { + "model": "gemini/gemini-2.5-flash", + "api_key": "test-key", + "order": 2, + "additional_drop_params": ["context_management"], + }, + }, + ], + ), + ( + "gemini-agent-deployment-id", + [ + { + "model_name": "gemini-agent", + "litellm_params": { + "model": "gemini/gemini-2.5-flash", + "api_key": "test-key", + "additional_drop_params": ["context_management"], + }, + "model_info": {"id": "gemini-agent-deployment-id"}, + }, + ], + ), + ( + "gemini-agent-deployment-id", + [ + {"model_name": "*", "litellm_params": {"model": "gemini/*", "api_key": "test-key"}}, + { + "model_name": "gemini-agent", + "litellm_params": { + "model": "gemini/gemini-2.5-flash", + "api_key": "test-key", + "additional_drop_params": ["context_management"], + }, + "model_info": {"id": "gemini-agent-deployment-id"}, + }, + ], + ), + ), + ids=("group-member-drops", "deployment-id-drops", "deployment-id-drops-beside-a-wildcard"), +) +async def test_gemini_messages_count_tokens_keeps_tool_results_when_a_routable_deployment_drops_context_management( + requested_model: str, model_list: list[dict[str, JsonValue]] +) -> None: + response, sent_bodies = await _count_through_the_messages_endpoint( + Router(model_list=model_list), {"model": requested_model} + ) + + assert response == {"input_tokens": 41} + assert sent_bodies[0].decode().count(_GEMINI_TOOL_OUTPUT) == 3 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("requested_model", "make_router", "request_options", "expected_tool_outputs_sent"), + ( + ( + "unrouted-model", + lambda: Router(model_list=[_gemini_deployment("gemini-agent", {})], default_fallbacks=["gemini-agent"]), + {}, + 1, + ), + ( + "unrouted-model", + lambda: Router( + model_list=[_gemini_deployment("gemini-agent", {"additional_drop_params": ["context_management"]})], + default_fallbacks=["gemini-agent"], + ), + {}, + 3, + ), + ( + "unrouted-model", + lambda: Router( + model_list=[_gemini_deployment("gemini-agent", {})], + default_fallbacks=["gemini-agent"], + default_litellm_params={"additional_drop_params": ["context_management"]}, + ), + {}, + 3, + ), + ( + "unrouted-model", + lambda: Router(model_list=[_gemini_deployment("gemini-agent", {})], default_fallbacks=["gemini-agent"]), + {"additional_drop_params": ["context_management"]}, + 3, + ), + ( + "flash-alias", + lambda: Router( + model_list=[_gemini_deployment("gemini-agent", {"additional_drop_params": ["context_management"]})], + model_group_alias={"flash-alias": "gemini/gemini-2.5-flash"}, + ), + {}, + 3, + ), + ), + ids=( + "default-fallback-edits-applied", + "default-fallback-deployment-drops-context-management", + "default-fallback-router-default-drops-context-management", + "default-fallback-request-drops-context-management", + "alias-to-a-litellm-model-deployment-drops-context-management", + ), +) +async def test_gemini_messages_count_tokens_reads_the_opt_out_of_the_deployment_the_router_falls_back_to( + requested_model: str, + make_router: Callable[[], Router], + request_options: dict[str, JsonValue], + expected_tool_outputs_sent: int, +) -> None: + response, sent_bodies = await _count_through_the_messages_endpoint( + make_router(), {**request_options, "model": requested_model} + ) + + assert response == {"input_tokens": 41} + assert sent_bodies[0].decode().count(_GEMINI_TOOL_OUTPUT) == expected_tool_outputs_sent + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("requested_model", "completion_model", "expected_tool_outputs_sent"), + (("gemini-agent", "gemini-dropping", 3), ("gemini-dropping", "gemini-agent", 1)), + ids=("completion-model-drops-context-management", "completion-model-keeps-context-management"), +) +async def test_gemini_messages_count_tokens_reads_the_opt_out_of_the_completion_model_inference_sends_to( + requested_model: str, completion_model: str, expected_tool_outputs_sent: int +) -> None: + import litellm.proxy.proxy_server as proxy_server + + router: Final = Router( + model_list=[ + _gemini_deployment("gemini-agent", {}), + _gemini_deployment("gemini-dropping", {"additional_drop_params": ["context_management"]}), + ] + ) + with patch.object(proxy_server, "general_settings", {"completion_model": completion_model}): + response, sent_bodies = await _count_through_the_messages_endpoint(router, {"model": requested_model}) + + assert response == {"input_tokens": 41} + assert sent_bodies[0].decode().count(_GEMINI_TOOL_OUTPUT) == expected_tool_outputs_sent + + class _GeminiPart(BaseModel): model_config = ConfigDict(frozen=True) text: str = "" @@ -1224,6 +1566,48 @@ def _gemini_router(drop_params: bool = False) -> Router: ) +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("deployment_params", "router_defaults", "request_options", "litellm_drop_params"), + ( + ({"drop_params": True}, {}, {}, False), + ({}, {"drop_params": True}, {}, False), + ({}, {}, {"drop_params": True}, False), + ({}, {}, {}, True), + ), + ids=("deployment-drop-params", "router-default-drop-params", "request-drop-params", "litellm-drop-params"), +) +async def test_gemini_messages_count_tokens_counts_the_unedited_body_when_drop_params_drops_an_invalid_edit( + deployment_params: dict[str, JsonValue], + router_defaults: dict[str, JsonValue], + request_options: dict[str, JsonValue], + litellm_drop_params: bool, +) -> None: + router: Final = Router( + model_list=[_gemini_deployment("gemini-agent", deployment_params)], default_litellm_params=router_defaults + ) + with patch.object(litellm, "drop_params", litellm_drop_params): + response, sent_bodies = await _count_through_the_messages_endpoint( + router, {**request_options, "model": "gemini-agent", "context_management": _INVALID_COMPACT_EDIT} + ) + + assert response == {"input_tokens": 41} + assert sent_bodies[0].decode().count(_GEMINI_TOOL_OUTPUT) == 3 + + +@pytest.mark.asyncio +async def test_gemini_messages_count_tokens_rejects_an_invalid_edit_without_drop_params() -> None: + with patch.object(litellm, "drop_params", False): + response, sent_bodies = await _count_through_the_messages_endpoint( + Router(model_list=[_gemini_deployment("gemini-agent", {})]), + {"model": "gemini-agent", "context_management": _INVALID_COMPACT_EDIT}, + ) + + assert isinstance(response, JSONResponse) + assert response.status_code == 400 + assert sent_bodies == () + + @pytest.mark.asyncio @pytest.mark.parametrize( "text", ["hello there", "The quick brown fox jumps over the lazy dog. " * 50], ids=["short", "long"] diff --git a/tests/unit/test_router/test_router_get_deployments.py b/tests/unit/test_router/test_router_get_deployments.py index 8b0c34f863c..07486762a64 100644 --- a/tests/unit/test_router/test_router_get_deployments.py +++ b/tests/unit/test_router/test_router_get_deployments.py @@ -5,6 +5,7 @@ import os import random import traceback from collections import defaultdict +from collections.abc import Callable from typing import Final, Literal import pytest @@ -110,9 +111,7 @@ def test_get_available_deployment_for_pass_through(): ) # Test that only pass-through deployment is returned - selected_model = router.get_available_deployment_for_pass_through( - "gpt-3.5-turbo" - ) + selected_model = router.get_available_deployment_for_pass_through("gpt-3.5-turbo") assert selected_model["litellm_params"]["model"] == "gpt-3.5-turbo" assert selected_model["litellm_params"]["use_in_pass_through"] is True @@ -259,9 +258,7 @@ def test_filter_pass_through_deployments(): healthy_deployments = router.get_model_list() # Filter pass-through deployments - pass_through_deployments = router._filter_pass_through_deployments( - healthy_deployments - ) + pass_through_deployments = router._filter_pass_through_deployments(healthy_deployments) # Should only have 2 deployments with use_in_pass_through=True assert len(pass_through_deployments) == 2 @@ -280,9 +277,7 @@ async def _weighted_async_deployment_ids(router: Router) -> tuple[str, ...]: random.seed(2025) deployments: Final = tuple( [ - await router.async_get_available_deployment( - model="shared", messages=None, request_kwargs={} - ) + await router.async_get_available_deployment(model="shared", messages=None, request_kwargs={}) for _ in range(1000) ] ) @@ -326,10 +321,7 @@ def test_weighted_selection_router( selected_ids: Final = ( asyncio.run(_weighted_async_deployment_ids(router)) if async_selection - else tuple( - router.get_available_deployment("shared")["model_info"]["id"] - for _ in range(1000) - ) + else tuple(router.get_available_deployment("shared")["model_info"]["id"] for _ in range(1000)) ) assert selected_ids.count("high") / len(selected_ids) > 0.89 @@ -355,10 +347,7 @@ def test_weighted_selection_router_no_rpm_set() -> None: }, ] ) - selected_ids: Final = tuple( - router.get_available_deployment("shared")["model_info"]["id"] - for _ in range(100) - ) + selected_ids: Final = tuple(router.get_available_deployment("shared")["model_info"]["id"] for _ in range(100)) assert set(selected_ids) == {"low", "high"} @@ -378,3 +367,61 @@ def test_model_group_aliases_select_the_target_group() -> None: assert deployment["model_name"] == "target" assert deployment["model_info"]["id"] == "target-deployment" + + +_ROUTABLE_MODEL_LIST: Final = [ + { + "model_name": "shared", + "litellm_params": {"model": "openai/shared-low", "api_key": "test-key"}, + "model_info": {"id": "low"}, + }, + { + "model_name": "shared", + "litellm_params": {"model": "openai/shared-high", "api_key": "test-key"}, + "model_info": {"id": "high"}, + }, + { + "model_name": "unrelated", + "litellm_params": {"model": "openai/unrelated", "api_key": "test-key"}, + "model_info": {"id": "unrelated"}, + }, +] + + +@pytest.mark.parametrize( + ("requested_model", "make_router", "expected_ids"), + ( + ("shared", lambda: Router(model_list=_ROUTABLE_MODEL_LIST), {"low", "high"}), + ("high", lambda: Router(model_list=_ROUTABLE_MODEL_LIST), {"high"}), + ("openai/unrelated", lambda: Router(model_list=_ROUTABLE_MODEL_LIST), {"unrelated"}), + ( + "unknown", + lambda: Router(model_list=_ROUTABLE_MODEL_LIST, default_fallbacks=["shared"]), + {"low", "high"}, + ), + ( + "alias", + lambda: Router(model_list=_ROUTABLE_MODEL_LIST, model_group_alias={"alias": "openai/unrelated"}), + {"unrelated"}, + ), + ("unknown", lambda: Router(model_list=_ROUTABLE_MODEL_LIST), set()), + ), + ids=( + "model-group", + "deployment-id", + "litellm-model-name", + "default-fallback", + "alias-to-a-litellm-model-name", + "unroutable", + ), +) +def test_get_routable_deployments_returns_every_deployment_routing_can_pick( + requested_model: str, make_router: Callable[[], Router], expected_ids: set[str] +) -> None: + router: Final = make_router() + + routable_ids: Final = { + deployment["model_info"]["id"] for deployment in router.get_routable_deployments(requested_model) + } + + assert routable_ids == expected_ids From f811a5923eb69684bd4b1deb8f93d0f87acc4efd Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 16 Oct 2026 11:32:34 -0700 Subject: [PATCH 09/58] test(cost): bill vertex gemini-3.8-flash reasoning tokens at the reasoning price (#46939) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../cost_calculation/cases/vertex_ai.py | 60 +------------------ .../integration/cost_calculation/cost_map.py | 1 + 2 files changed, 3 insertions(+), 58 deletions(-) diff --git a/tests/integration/cost_calculation/cases/vertex_ai.py b/tests/integration/cost_calculation/cases/vertex_ai.py index e35fd2307c2..e19e377ae93 100644 --- a/tests/integration/cost_calculation/cases/vertex_ai.py +++ b/tests/integration/cost_calculation/cases/vertex_ai.py @@ -2181,61 +2181,6 @@ GEMINI_3_8_FLASH_GOOGLE_MAPS_GROUNDING: Final = CostTrackingTestCase( ), ) -GEMINI_3_8_FLASH_FALLBACK_REASONING_AT_OUTPUT_RATE: Final = CostTrackingTestCase( - name="gemini-3.8-flash-fallback_reasoning_at_output_rate", - covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="gemini-3.8-flash", - request={ - "model": "$MODEL", - "messages": [ - { - "role": "system", - "content": [ - { - "type": "text", - "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", - } - ], - }, - { - "role": "user", - "content": [ - { - "type": "text", - "text": "f92946792f44 summarize the attached material in one line and name the city weather", - } - ], - }, - ], - "stream": False, - "reasoning_effort": "medium", - "allowed_openai_params": ["reasoning_effort"], - }, - response=JsonResponse( - content_type="application/json", - body={ - "candidates": [ - { - "content": {"parts": [{"text": "scripted answer f92946792f44"}], "role": "model"}, - "finishReason": "STOP", - "index": 0, - } - ], - "usageMetadata": { - "promptTokenCount": 1240, - "candidatesTokenCount": 560, - "thoughtsTokenCount": 3480, - "totalTokenCount": 5280, - "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 1240}], - }, - "modelVersion": "gemini-3.8-flash", - }, - ), - expected=ExactExpected( - spend=0.0132496, input_cost=0.0006448, output_cost=0.0126048, prompt_tokens=1240, completion_tokens=4040 - ), -) - GEMINI_3_8_FLASH_FALLBACK_IMAGE_TOKENS_AT_INPUT_RATE: Final = CostTrackingTestCase( name="gemini-3.8-flash-fallback_image_tokens_at_input_rate", covers="quota_management.spend_tracking.cost_matrix.logs_cost", @@ -3020,7 +2965,7 @@ GEMINI_3_8_FLASH_STREAM_FULL_USAGE: Final = CostTrackingTestCase( ), ), expected=ExactExpected( - spend=0.007460128, input_cost=0.001619488, output_cost=0.00584064, prompt_tokens=8314, completion_tokens=1592 + spend=0.006992128, input_cost=0.001619488, output_cost=0.00537264, prompt_tokens=8314, completion_tokens=1592 ), ) @@ -3368,7 +3313,7 @@ GEMINI_3_1_PRO_PREVIEW_CACHE_READ: Final = _gemini_case( ExactExpected(spend=0.00912736, input_cost=0.00411136, output_cost=0.005016, prompt_tokens=12928, completion_tokens=380) ) GEMINI_3_1_PRO_PREVIEW_REASONING: Final = _gemini_case( - GEMINI_3_8_FLASH_FALLBACK_REASONING_AT_OUTPUT_RATE, "gemini-3.1-pro-preview-reasoning", "gemini-3.1-pro-preview", + GEMINI_3_1_PRO_REASONING, "gemini-3.1-pro-preview-reasoning", "gemini-3.1-pro-preview", ExactExpected(spend=0.056056, input_cost=0.002728, output_cost=0.053328, prompt_tokens=1240, completion_tokens=4040) ) GEMINI_3_1_PRO_PREVIEW_TOOL_CALL: Final = _gemini_case( @@ -3410,7 +3355,6 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GEMINI_3_8_FLASH_SERVICE_TIER_PRIORITY, GEMINI_3_8_FLASH_WEB_SEARCH_PER_PROMPT, GEMINI_3_8_FLASH_GOOGLE_MAPS_GROUNDING, - GEMINI_3_8_FLASH_FALLBACK_REASONING_AT_OUTPUT_RATE, GEMINI_3_8_FLASH_FALLBACK_IMAGE_TOKENS_AT_INPUT_RATE, GEMINI_3_8_FLASH_STREAM, GEMINI_3_8_FLASH_STREAM_NO_USAGE, diff --git a/tests/integration/cost_calculation/cost_map.py b/tests/integration/cost_calculation/cost_map.py index 40072bfcaf6..686636dacbd 100644 --- a/tests/integration/cost_calculation/cost_map.py +++ b/tests/integration/cost_calculation/cost_map.py @@ -579,6 +579,7 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "mode": "chat", "output_cost_per_audio_token": 6.24e-06, "output_cost_per_token": 3.12e-06, + "output_cost_per_reasoning_token": 2.6e-06, "output_cost_per_token_flex": 1.56e-06, "output_cost_per_token_priority": 3.9e-06, "search_context_cost_per_query": {"search_context_size_medium": 0.035}, From 339fd90836c1ce22dd18d519fd0c4af5f66544ce Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Fri, 16 Oct 2026 11:33:01 -0700 Subject: [PATCH 10/58] fix(proxy): streaming responses hand the event loop back every 50 ms instead of holding it for the whole buffered stream (#46925) * fix(proxy): streaming responses hand the event loop back every 50 ms instead of holding it for the whole buffered stream * fix(proxy): yield to the event loop on every streamed chunk, no hold budget or clock * docs(proxy): say why the streaming loop yields on every chunk * fix(proxy): streaming loop yields once it has held the event loop 50 ms, clock read inline * docs(proxy): streaming loop comment names the 50 ms hold * Revert "docs(proxy): streaming loop comment names the 50 ms hold" This reverts commit cbf1371873c0540b83fdcf4e9a2bafb69b72f460. * refactor(proxy): streaming loop hold budget is _STREAM_LOOP_HOLD_MS = 50 * Revert "refactor(proxy): streaming loop hold budget is _STREAM_LOOP_HOLD_MS = 50" This reverts commit 356454def6b5efed7a6afea69f5f059468227ea7. * docs(proxy): streaming loop comment notes the yield is a floor, not a count of releases --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 8 +++++ .../proxy_server/test_streaming_helpers.py | 33 +++++++++++++++++++ 2 files changed, 41 insertions(+) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 89fd3ae0f49..3de3dccc532 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -10013,6 +10013,7 @@ def _pop_complete_sse_frame(buffer: str) -> tuple[str | None, str]: _STREAM_KEEPALIVE: Final = object() _KEEPALIVE_MIN_SECONDS: Final = 1.0 +_STREAM_LOOP_HOLD_SECONDS: Final = 0.05 _KEEPALIVE_MAX_SECONDS: Final = 300.0 _EMPTY_MAPPING: Final[Mapping[str, object]] = MappingProxyType({}) _EMPTY_HEADERS: Final[Mapping[str, str]] = MappingProxyType({}) @@ -10288,7 +10289,14 @@ async def async_data_generator( else stream_iterator ) + loop_released_at = time.monotonic() # rebind-ok: last time this stream let other tasks run async for item in stream_source: + # Unblock the event loop: a stream whose chunks are already buffered never awaits + # between chunks, so without this yield one long stream holds the loop for its whole body. + # A real await in between may already have released the loop; this only guarantees it + if time.monotonic() - loop_released_at >= _STREAM_LOOP_HOLD_SECONDS: + await asyncio.sleep(0) + loop_released_at = time.monotonic() if item is _STREAM_KEEPALIVE: yield ": ping\n\n" continue diff --git a/tests/unit/proxy/proxy_server/test_streaming_helpers.py b/tests/unit/proxy/proxy_server/test_streaming_helpers.py index 09d71cb4ead..f53e1931782 100644 --- a/tests/unit/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/unit/proxy/proxy_server/test_streaming_helpers.py @@ -2090,3 +2090,36 @@ def test_fast_serialize_simple_model_response_stream_keeps_served_service_tier() assert result is not None assert json.loads(result)["service_tier"] == "priority" + + +@pytest.mark.asyncio +async def test_async_data_generator_lets_other_tasks_run_during_a_stream_that_never_waits(monkeypatch): + """A stream whose chunks are already buffered never suspends on its own, so + the generator has to hand the loop back itself once it has held the loop + long enough: another task must get to run before the stream is over, and + not before the hold budget is used up.""" + _patch_logging_flags(monkeypatch) + order: list[str] = [] # mutable-ok: records the interleaving of the two tasks + clock: list[float] = [0.0] # mutable-ok: fake monotonic clock, advanced 20 ms per chunk + monkeypatch.setattr(ps.time, "monotonic", lambda: clock[0]) + + async def chunks(): + for index in range(6): + clock[0] += 0.02 + order.append(f"chunk{index}") + yield _simple_chunk(content=str(index)) + + async def other_task(): + order.append("other") + + other = asyncio.ensure_future(other_task()) + async for _ in async_data_generator( + response=chunks(), + user_api_key_dict=_user_auth(), + request_data={"model": "gpt-4"}, + ): + pass + await other + + assert "other" in order[:-1], order + assert order.index("other") >= 3, order From 1dfceccf10fcf23deeae69d361390e548cfb03f3 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Fri, 16 Oct 2026 11:50:27 -0700 Subject: [PATCH 11/58] test(perf): cap every awaited Postgres statement at 0.5 calls per request (#46944) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../perf/cases/mixed_size_requests_2rps.py | 1 + tests/e2e/perf/cases/small_bodies_5rps.py | 1 + tests/e2e/perf/setup/observe.py | 24 ++++++++- tests/e2e/perf/setup/performance_test_case.py | 1 + tests/e2e_harness/perf/setup/test_observe.py | 54 +++++++++++++++++++ 5 files changed, 80 insertions(+), 1 deletion(-) diff --git a/tests/e2e/perf/cases/mixed_size_requests_2rps.py b/tests/e2e/perf/cases/mixed_size_requests_2rps.py index 41b61cdc987..f95fb0845c4 100644 --- a/tests/e2e/perf/cases/mixed_size_requests_2rps.py +++ b/tests/e2e/perf/cases/mixed_size_requests_2rps.py @@ -69,6 +69,7 @@ CASE: Final = PerformanceTestCase( Limit("memory_growth_megabytes", max=600), Limit("cpu_seconds_per_request", max=1.26), Limit("postgres_calls_per_request", max=3.5), + Limit("postgres_max_statement_per_request", max=0.5), Limit("redis_calls_per_request", max=21.0), ), ) diff --git a/tests/e2e/perf/cases/small_bodies_5rps.py b/tests/e2e/perf/cases/small_bodies_5rps.py index 00a9fdefdb5..9d3484de05d 100644 --- a/tests/e2e/perf/cases/small_bodies_5rps.py +++ b/tests/e2e/perf/cases/small_bodies_5rps.py @@ -68,6 +68,7 @@ CASE: Final = PerformanceTestCase( Limit("time_to_first_token_seconds", percentile=99, max=3.7), Limit("memory_growth_megabytes", max=100), Limit("postgres_calls_per_request", max=1.8), + Limit("postgres_max_statement_per_request", max=0.5), Limit("redis_calls_per_request", max=18), Limit("cpu_seconds_per_request", max=0.29), ), diff --git a/tests/e2e/perf/setup/observe.py b/tests/e2e/perf/setup/observe.py index 98391a02418..be64679c9da 100644 --- a/tests/e2e/perf/setup/observe.py +++ b/tests/e2e/perf/setup/observe.py @@ -128,6 +128,13 @@ class Observation: def statements(self) -> tuple[StatementCalls, ...] | None: return statement_deltas(self.db_before, self.db_after) + @property + def busiest_awaited_statement(self) -> StatementCalls | None: + """The awaited statement with the most calls: a lookup added to every request shows up here at + 1.0 per request while cache refreshes stay fractional, which the awaited total cannot tell apart.""" + awaited: Final = tuple(item for item in self.statements or () if item.bucket == "awaited") + return max(awaited, key=_item_count, default=None) + @property def retried_markers(self) -> frozenset[str]: counts: Final = Counter(seen.marker for seen in self.mock_requests if seen.marker is not None) @@ -209,6 +216,11 @@ def observe(metric: Metric, percentile_rank: int | None, observation: Observatio return observation.count_per_request(calls_in("awaited", observation.statements)) if metric == "postgres_background_calls_per_request": return observation.count_per_request(calls_in("background", observation.statements)) + if metric == "postgres_max_statement_per_request": + if observation.statements is None: + return None + busiest: Final = observation.busiest_awaited_statement + return observation.count_per_request(0 if busiest is None else busiest.calls) if metric == "redis_calls_per_request": return observation.count_per_request(calls_in("awaited", observation.redis_commands)) if metric == "redis_background_calls_per_request": @@ -225,6 +237,7 @@ class Row: percentile: int | None observed: float | None limit: float | None + detail: str = "" @property def label(self) -> str: @@ -252,6 +265,7 @@ REPORTED: Final[tuple[tuple[Metric, int | None], ...]] = ( ("memory_growth_megabytes", None), ("postgres_calls_per_request", None), ("postgres_background_calls_per_request", None), + ("postgres_max_statement_per_request", None), ("redis_calls_per_request", None), ("redis_background_calls_per_request", None), ("cpu_seconds_per_request", None), @@ -275,10 +289,11 @@ def evaluate(case: PerformanceTestCase, observation: Observation) -> tuple[list[ percentile=rank, observed=observe(metric, rank, observation), limit=limit.max if limit else None, + detail=_detail(metric, observation), ) ) exceeded: Final = [ - f"{row.label}: observed {row.observed:.4f} > limit {row.limit:.4f}" + f"{row.label}: observed {row.observed:.4f} > limit {row.limit:.4f}{f' ({row.detail})' if row.detail else ''}" if row.over and row.observed is not None and row.limit is not None else f"{row.label}: limit set but nothing observed" for row in rows @@ -287,6 +302,13 @@ def evaluate(case: PerformanceTestCase, observation: Observation) -> tuple[list[ return rows, exceeded +def _detail(metric: Metric, observation: Observation) -> str: + if metric != "postgres_max_statement_per_request": + return "" + busiest: Final = observation.busiest_awaited_statement + return "" if busiest is None else busiest.summary + + def _row_properties(row: Row) -> tuple[tuple[str, float], ...]: observed: Final = ((row.property_name, row.observed),) if row.observed is not None else () limit: Final = ((f"{row.property_name}_limit", row.limit),) if row.limit is not None else () diff --git a/tests/e2e/perf/setup/performance_test_case.py b/tests/e2e/perf/setup/performance_test_case.py index 76e6efa3203..6a734e2923f 100644 --- a/tests/e2e/perf/setup/performance_test_case.py +++ b/tests/e2e/perf/setup/performance_test_case.py @@ -23,6 +23,7 @@ Metric = Literal[ "memory_growth_megabytes", "postgres_calls_per_request", "postgres_background_calls_per_request", + "postgres_max_statement_per_request", "cpu_seconds_per_request", "redis_calls_per_request", "redis_background_calls_per_request", diff --git a/tests/e2e_harness/perf/setup/test_observe.py b/tests/e2e_harness/perf/setup/test_observe.py index f38e7927fb7..4b7e2d64a14 100644 --- a/tests/e2e_harness/perf/setup/test_observe.py +++ b/tests/e2e_harness/perf/setup/test_observe.py @@ -69,5 +69,59 @@ def test_junit_properties_carry_each_observed_metric_and_each_limit_including_un "perf_gateway_added_seconds_p99": 0.75, "perf_postgres_calls_per_request": 10.0, "perf_postgres_background_calls_per_request": 0.0, + "perf_postgres_max_statement_per_request": 10.0, "perf_redis_calls_per_request_limit": 35.0, } + + +def _statement(queryid: int, query: str, calls: int) -> StatementCalls: + return StatementCalls(queryid=queryid, query=query, calls=calls) + + +def _observation_with(statements: tuple[StatementCalls, ...]) -> Observation: + before: Final = tuple(replace(item, calls=0) for item in statements) + return Observation( + records=tuple(served(seq, 1.0) for seq in range(1000)), + mock_requests=(), + guardrail_calls=(), + rss_before=NO_RSS, + rss_after=NO_RSS, + wall_seconds=600.0, + db_before=DbCounters(commits=0, statements=before), + db_after=DbCounters(commits=0, statements=statements), + ) + + +MAX_STATEMENT_CASE: Final = replace(CASE, limits=(Limit("postgres_max_statement_per_request", max=0.5),)) +TEAM_REFRESH: Final = 'SELECT * FROM "LiteLLM_TeamMembership" WHERE user_id = $1' +USER_LOOKUP: Final = 'SELECT * FROM "LiteLLM_UserTable" WHERE user_id = $1' + + +def test_cache_refreshes_whose_sum_exceeds_the_cap_pass_the_per_statement_limit() -> None: + statements: Final = ( + _statement(1, TEAM_REFRESH, calls=357), + _statement(2, USER_LOOKUP, calls=245), + _statement(3, 'INSERT INTO "LiteLLM_DailyModelUsage" VALUES ($1)', calls=1000), + ) + + rows, exceeded = evaluate(MAX_STATEMENT_CASE, _observation_with(statements)) + + assert exceeded == [] + assert dict(junit_properties(rows))["perf_postgres_max_statement_per_request"] == 0.357 + assert dict(junit_properties(rows))["perf_postgres_calls_per_request"] == 0.602 + + +def test_a_lookup_on_every_request_fails_the_per_statement_limit_and_names_the_statement() -> None: + statements: Final = (_statement(1, TEAM_REFRESH, calls=357), _statement(2, USER_LOOKUP, calls=1000)) + + _, exceeded = evaluate(MAX_STATEMENT_CASE, _observation_with(statements)) + + assert exceeded == [f"postgres_max_statement_per_request: observed 1.0000 > limit 0.5000 ({USER_LOOKUP})"], exceeded + + +def test_per_statement_limit_without_db_counters_reports_nothing_observed() -> None: + observation: Final = replace(_observation_with(()), db_before=None, db_after=None) + + _, exceeded = evaluate(MAX_STATEMENT_CASE, observation) + + assert exceeded == ["postgres_max_statement_per_request: limit set but nothing observed"] From bca7650e8f6612375496797c805eb464cb2633bc Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Fri, 16 Oct 2026 12:16:55 -0700 Subject: [PATCH 12/58] feat(mcp): accept an ID-JAG as an RFC 7523 jwt-bearer grant on the gateway token endpoint (#46748) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/auth/user_api_key_auth_mcp.py | 4 +- .../mcp_server/discoverable_endpoints.py | 56 +- .../mcp_server/gateway_dcr_flow.py | 162 +++- .../_experimental/mcp_server/id_jag_grant.py | 540 +++++++++++++ .../mcp_server/idp_token_exchange.py | 41 +- .../_experimental/mcp_server/oauth_utils.py | 5 + .../mcp_server/proxy_api_credentials.py | 18 +- litellm/proxy/_lazy_features.py | 19 +- litellm/proxy/_lazy_openapi_snapshot.json | 86 +- litellm/proxy/_types.py | 25 + litellm/proxy/auth/auth_checks.py | 4 + litellm/proxy/auth/handle_jwt.py | 96 ++- litellm/proxy/auth/route_checks.py | 11 +- .../auth/test_user_api_key_auth_mcp.py | 117 +++ .../mcp_server/test_discoverable_endpoints.py | 218 +++++- .../mcp_server/test_gateway_dcr_flow.py | 204 ++++- .../mcp_server/test_id_jag_grant.py | 738 ++++++++++++++++++ .../mcp_server/test_proxy_api_credentials.py | 89 ++- tests/unit/proxy/auth/test_handle_jwt.py | 198 +++++ tests/unit/proxy/test__lazy_features.py | 22 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 8 +- 21 files changed, 2535 insertions(+), 126 deletions(-) create mode 100644 litellm/proxy/_experimental/mcp_server/id_jag_grant.py create mode 100644 tests/unit/proxy/_experimental/mcp_server/test_id_jag_grant.py diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 9c8890efd2b..4090666e16f 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -1345,7 +1345,9 @@ class MCPRequestHandler: if auth is None: return None refreshed = auth.model_copy() - if auth.mcp_admitted_user_subject and auth.user_id: + if (auth.mcp_admitted_user_subject or auth.is_session_token) and auth.user_id: + # A session credential (``lite login`` or a redeemed ID-JAG) has no key row, its authority is + # sealed in the token: refresh the live user policy and keep the token's team, routes and models. current = await MCPRequestHandler.reload_admitted_user(auth.user_id, requires_fresh_policy=True) refreshed.object_permission = current.object_permission refreshed.object_permission_id = current.object_permission_id diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 6e5b5486696..449adaa73e9 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -56,6 +56,7 @@ from litellm.proxy._experimental.mcp_server.faults import ( render_token_fault, ) from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( + JWT_BEARER_GRANT_TYPE, VendorCredentialState, aggregate_authorize, aggregate_token, @@ -71,6 +72,11 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( revoke_refresh_token, supported_grant_types, ) +from litellm.proxy._experimental.mcp_server.id_jag_grant import ( + basic_client_id, + jwt_bearer_available, + redeem_id_jag_assertion, +) from litellm.proxy._experimental.mcp_server.idp_token_exchange import ( exchange_idp_subject_token, token_exchange_available, @@ -2396,7 +2402,7 @@ async def token_endpoint( grant_type: str = Form(...), code: str = Form(None), redirect_uri: str = Form(None), - client_id: str = Form(...), + client_id: str | None = Form(None), client_secret: str | None = Form(None), code_verifier: str = Form(None), refresh_token: str | None = Form(None), @@ -2405,6 +2411,7 @@ async def token_endpoint( subject_token: str | None = Form(None), subject_token_type: str | None = Form(None), requested_token_type: str | None = Form(None), + assertion: str | None = Form(None), mcp_server_name: str | None = None, ): """ @@ -2416,7 +2423,14 @@ async def token_endpoint( 3. Return the token 4. Return a virtual key in this response """ - if mcp_server_name is None and is_gateway_dcr_client_id(client_id): + presented_client_id: Final = client_id or basic_client_id(request.headers.get("authorization")) + if presented_client_id is None: + raise HTTPException( + status_code=400, detail={"error": "invalid_request", "error_description": "client_id is required"} + ) + if mcp_server_name is None and ( + is_gateway_dcr_client_id(presented_client_id) or grant_type == JWT_BEARER_GRANT_TYPE + ): from litellm.proxy.proxy_server import ( # noqa: PLC0415 # circular import at module load master_key, user_api_key_cache, @@ -2427,7 +2441,7 @@ async def token_endpoint( grant_type=grant_type, code=code, redirect_uri=redirect_uri, - client_id=client_id, + client_id=presented_client_id, code_verifier=code_verifier, refresh_token=refresh_token, master_key=master_key, @@ -2439,9 +2453,12 @@ async def token_endpoint( subject_token_type=subject_token_type, requested_token_type=requested_token_type, exchange_subject_token=exchange_idp_subject_token, + assertion=assertion, + client_secret=client_secret, + redeem_assertion=redeem_id_jag_assertion, ) - lookup_name: Final = mcp_server_name or client_id + lookup_name: Final = mcp_server_name or presented_client_id client_ip: Final = IPAddressUtils.get_mcp_client_ip(request) mcp_server = _resolve_mcp_server_by_name_or_id(lookup_name, client_ip) if mcp_server is None and mcp_server_name is None: @@ -2454,7 +2471,7 @@ async def token_endpoint( grant_type=grant_type, code=code, redirect_uri=redirect_uri, - client_id=client_id, + client_id=presented_client_id, client_secret=client_secret, code_verifier=code_verifier, refresh_token=refresh_token, @@ -2568,7 +2585,8 @@ async def native_client_auth_discovery(request: Request) -> JSONResponse: """The versioned contract a native client (``lite login --pkce``, or a CLI in any other language) reads to sign a user in through the browser and obtain a proxy credential.""" return JSONResponse( - native_client_auth_contract(request, token_exchange_available()), headers=TOKEN_NO_CACHE_HEADERS + native_client_auth_contract(request, token_exchange_available(), jwt_bearer_available()), + headers=TOKEN_NO_CACHE_HEADERS, ) @@ -3135,7 +3153,7 @@ def _aggregate_header_scopes(request: Request) -> tuple[str, ...]: def _build_aggregate_authorization_server_response( - request: Request, token_exchange_available: bool + request: Request, token_exchange_available: bool, jwt_bearer_available: bool = False ) -> dict[str, object]: """RFC 8414 metadata for the gateway as the aggregate authorization server. @@ -3155,9 +3173,13 @@ def _build_aggregate_authorization_server_response( "registration_endpoint": f"{request_base_url}/register", "response_types_supported": ["code"], "scopes_supported": [], - "grant_types_supported": supported_grant_types(token_exchange_available), + "grant_types_supported": supported_grant_types(token_exchange_available, jwt_bearer_available), "code_challenge_methods_supported": ["S256"], - "token_endpoint_auth_methods_supported": ["none", "client_secret_post"], + "token_endpoint_auth_methods_supported": ( + ["none", "client_secret_post", "client_secret_basic"] + if jwt_bearer_available + else ["none", "client_secret_post"] + ), } @@ -3193,7 +3215,7 @@ async def oauth_authorization_server_aggregate(request: Request): per-server row win here instead would serve an issuer of {base} against a resource that advertised {base}/mcp, which fails the RFC 8414 issuer check and breaks the front door. """ - return _build_aggregate_authorization_server_response(request, token_exchange_available()) + return _build_aggregate_authorization_server_response(request, token_exchange_available(), jwt_bearer_available()) # Standard MCP pattern: /.well-known/oauth-protected-resource/mcp/{server_name} @@ -3273,6 +3295,10 @@ def _build_oauth_authorization_server_response( if explicitly_named else request_base_url ) + # With no upstream resolved the document is the gateway's own: its issuer is the bare origin + # and its /token serves every grant the gateway does, including the ID-JAG jwt-bearer grant. + gateway_owned: Final = mcp_server_name is None and issuer_path is None + jwt_bearer: Final = gateway_owned and jwt_bearer_available() return { "issuer": issuer, @@ -3280,9 +3306,15 @@ def _build_oauth_authorization_server_response( "token_endpoint": token_endpoint, "response_types_supported": ["code"], "scopes_supported": (mcp_server.scopes if mcp_server and mcp_server.scopes else []), - "grant_types_supported": ["authorization_code", "refresh_token"], + "grant_types_supported": ( + list(supported_grant_types(token_exchange_available(), jwt_bearer)) + if gateway_owned + else ["authorization_code", "refresh_token"] + ), "code_challenge_methods_supported": ["S256"], - "token_endpoint_auth_methods_supported": ["client_secret_post"], + "token_endpoint_auth_methods_supported": ( + ["none", "client_secret_post", "client_secret_basic"] if jwt_bearer else ["client_secret_post"] + ), # Claude expects a registration endpoint, even if we just fake it "registration_endpoint": ( f"{request_base_url}/{mcp_server_name}/register" if mcp_server_name else f"{request_base_url}/register" diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index dd8c82a99fe..90289fd4c20 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -198,20 +198,26 @@ class MintProxyCredential(Protocol): and mints the same credential ``lite login`` mints.""" def __call__( - self, user_id: str, team_id: str | None, / + self, user_id: str, team_id: str | None, /, *, allowed_routes: Sequence[str] | None = None ) -> Awaitable[MintedProxyCredential | ProxyCredentialMintFailure]: ... TOKEN_EXCHANGE_GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:token-exchange" +JWT_BEARER_GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:jwt-bearer" -def supported_grant_types(token_exchange_available: bool) -> tuple[str, ...]: +def supported_grant_types(token_exchange_available: bool, jwt_bearer_available: bool = False) -> tuple[str, ...]: """The grants ``/token`` can serve on this deployment. The RFC 8693 exchange is listed only where the JWT auth that proves a subject token is on, backed by a database, and - licensed, so a client never selects a grant the gateway would then refuse.""" - if token_exchange_available: - return ("authorization_code", "refresh_token", TOKEN_EXCHANGE_GRANT_TYPE) - return ("authorization_code", "refresh_token") + licensed, so a client never selects a grant the gateway would then refuse. The RFC 7523 + jwt-bearer grant (an ID-JAG from the identity provider) is listed only where, on top of + that, the operator has allowlisted at least one requesting app.""" + return ( + "authorization_code", + "refresh_token", + *((TOKEN_EXCHANGE_GRANT_TYPE,) if token_exchange_available else ()), + *((JWT_BEARER_GRANT_TYPE,) if jwt_bearer_available else ()), + ) """RFC 8693: a native client that already holds a token from the customer's identity @@ -236,10 +242,29 @@ class SubjectIdentity(LiteLLMBaseModel): class SubjectTokenRefusal(LiteLLMBaseModel): model_config = ConfigDict(frozen=True) - error: Literal["unsupported_grant_type", "invalid_request", "temporarily_unavailable"] + error: Literal[ + "unsupported_grant_type", + "invalid_request", + "invalid_grant", + "invalid_client", + "invalid_scope", + "temporarily_unavailable", + ] description: str = Field(min_length=1) +class AssertionIdentity(SubjectIdentity): + """The litellm user an RFC 7523 assertion (an ID-JAG) stands for, plus what the grant + tail needs from the assertion itself: the single-use key for its ``jti``, how long that + marker must outlive it, the ``scope`` the credential is issued for, and the route groups that + scope narrows the credential to (``None`` leaves it the full ``lite login`` credential).""" + + replay_key: str = Field(min_length=1) + replay_ttl_seconds: int = Field(gt=0) + scope: str | None = None + allowed_routes: tuple[str, ...] | None = None + + class ExchangeSubjectToken(Protocol): """Injected RFC 8693 subject-token verifier ``(subject_token, request)``: proves the IdP token the way the proxy's own JWT auth does and names the litellm user and team it @@ -272,7 +297,9 @@ class LookupConsentTeams(Protocol): def __call__(self, user_id: str, /) -> Awaitable[tuple[ConsentTeam, ...] | ReloadUserFailure]: ... -async def _refuse_proxy_credential(user_id: str, team_id: str | None) -> ProxyCredentialMintFailure: +async def _refuse_proxy_credential( + user_id: str, team_id: str | None, *, allowed_routes: Sequence[str] | None = None +) -> ProxyCredentialMintFailure: return "unresolvable" @@ -282,6 +309,25 @@ async def _refuse_subject_token(subject_token: str, request: Request) -> Subject ) +class RedeemAssertion(Protocol): + """Injected RFC 7523 assertion verifier ``(assertion, client_id, client_secret, request)``: + proves an ID-JAG the identity provider minted for the requesting app ``client_id``, + authenticates that app, and names the litellm user the assertion stands for, or says why + this gateway will not take it.""" + + def __call__( + self, assertion: str, client_id: str, client_secret: str | None, request: Request, / + ) -> Awaitable[AssertionIdentity | SubjectTokenRefusal]: ... + + +async def _refuse_assertion( + assertion: str, client_id: str, client_secret: str | None, request: Request +) -> SubjectTokenRefusal: + return SubjectTokenRefusal( + error="unsupported_grant_type", description="this gateway is not configured to redeem identity assertions" + ) + + async def _unavailable_vendor_credential( user_id: str, server_id: str, requested_scopes: tuple[str, ...] = () ) -> VendorCredentialState: @@ -797,7 +843,9 @@ class NativeClientAuthContract(TypedDict): revocation_endpoint_auth_methods_supported: ReadOnly[tuple[str, ...]] -def native_client_auth_contract(request: Request, token_exchange_available: bool) -> NativeClientAuthContract: +def native_client_auth_contract( + request: Request, token_exchange_available: bool, jwt_bearer_available: bool = False +) -> NativeClientAuthContract: """The versioned discovery document at ``/.well-known/litellm-cli-auth``: everything a native client (in any language) needs to run the sign-in without reading LiteLLM source. ``resource`` is the exact value to send as the RFC 8707 ``resource`` parameter @@ -812,7 +860,7 @@ def native_client_auth_contract(request: Request, token_exchange_available: bool "revocation_endpoint": f"{base_url}/revoke", "resource": base_url, "response_types_supported": ("code",), - "grant_types_supported": supported_grant_types(token_exchange_available), + "grant_types_supported": supported_grant_types(token_exchange_available, jwt_bearer_available), "code_challenge_methods_supported": ("S256",), "token_endpoint_auth_methods_supported": ("none",), "revocation_endpoint_auth_methods_supported": ("none",), @@ -1403,10 +1451,11 @@ class _ProxyCredentialTokenResponse(TypedDict): access_token: ReadOnly[str] token_type: ReadOnly[Literal["Bearer"]] expires_in: ReadOnly[int] - refresh_token: ReadOnly[str] + refresh_token: NotRequired[ReadOnly[str]] user_id: ReadOnly[str] team_id: ReadOnly[str | None] issued_token_type: NotRequired[ReadOnly[_IssuedTokenType]] + scope: NotRequired[ReadOnly[str]] def proxy_credential_response( @@ -1415,6 +1464,7 @@ def proxy_credential_response( signing: SessionSigning, issued_token_type: _IssuedTokenType | None = None, family: str | None = None, + scope: str | None = None, ) -> Response: """The proxy-API token response: the access token is the very credential ``lite login`` stores (accepted on every proxy route with user and team attribution), and @@ -1426,7 +1476,7 @@ def proxy_credential_response( refresh: Final = mint_session_refresh_token(bound_principal, signing.keys, signing.now, family=family) if not isinstance(refresh, MintedSessionToken): return _oauth_error(500, "server_error", "failed to mint the session credential") - credential: Final[_ProxyCredentialTokenResponse] = { + renewable: Final[_ProxyCredentialTokenResponse] = { "access_token": minted.key, "token_type": "Bearer", "expires_in": minted.expires_in, @@ -1434,12 +1484,27 @@ def proxy_credential_response( "user_id": minted.user_id, "team_id": minted.team_id, } + credential: Final[_ProxyCredentialTokenResponse] = renewable if scope is None else {**renewable, "scope": scope} if issued_token_type is None: return JSONResponse(status_code=200, content=credential, headers=TOKEN_NO_CACHE_HEADERS) exchanged: Final[_ProxyCredentialTokenResponse] = {**credential, "issued_token_type": issued_token_type} return JSONResponse(status_code=200, content=exchanged, headers=TOKEN_NO_CACHE_HEADERS) +def assertion_credential_response(minted: MintedProxyCredential, scope: str | None) -> Response: + """The proxy-API token response for an assertion grant: the same access token, with the + scope the assertion carried, and no refresh token.""" + unscoped: Final[_ProxyCredentialTokenResponse] = { + "access_token": minted.key, + "token_type": "Bearer", + "expires_in": minted.expires_in, + "user_id": minted.user_id, + "team_id": minted.team_id, + } + credential: Final[_ProxyCredentialTokenResponse] = unscoped if scope is None else {**unscoped, "scope": scope} + return JSONResponse(status_code=200, content=credential, headers=TOKEN_NO_CACHE_HEADERS) + + def _reload_failure_response(failure: ReloadUserFailure) -> Response: """Map the live-user revalidation failure onto its OAuth error, exhaustively, so a new ``ReloadUserFailure`` member is a type error here rather than silently 400ing.""" @@ -1460,7 +1525,9 @@ def _subject_token_refusal_response(refusal: SubjectTokenRefusal) -> Response: match refusal.error: case "temporarily_unavailable": return _oauth_error(503, refusal.error, refusal.description) - case "unsupported_grant_type" | "invalid_request": + case "invalid_client": + return _oauth_error(401, refusal.error, refusal.description) + case "unsupported_grant_type" | "invalid_request" | "invalid_grant" | "invalid_scope": return _oauth_error(400, refusal.error, refusal.description) case _: assert_never(refusal.error) @@ -1515,6 +1582,9 @@ async def aggregate_token( subject_token_type: str | None = None, requested_token_type: str | None = None, exchange_subject_token: ExchangeSubjectToken = _refuse_subject_token, + assertion: str | None = None, + client_secret: str | None = None, + redeem_assertion: RedeemAssertion = _refuse_assertion, ) -> Response: """The aggregate token verb: authorization_code and refresh_token grants for the identity-only session pair, or for the proxy-API credential when the grant was issued @@ -1561,10 +1631,19 @@ async def aggregate_token( exchange_subject_token=exchange_subject_token, issue=issue, ) + if grant_type == JWT_BEARER_GRANT_TYPE: + return await _jwt_bearer_grant( + request=request, + assertion=assertion, + client_id=client_id, + client_secret=client_secret, + redeem_assertion=redeem_assertion, + issue=issue, + ) return _oauth_error( 400, "unsupported_grant_type", - f"grant_type must be authorization_code, refresh_token, or {TOKEN_EXCHANGE_GRANT_TYPE}", + f"grant_type must be authorization_code, refresh_token, {TOKEN_EXCHANGE_GRANT_TYPE}, or {JWT_BEARER_GRANT_TYPE}", ) @@ -1690,6 +1769,48 @@ class _GrantIssuer: return _mint_failure_response(minted) return proxy_credential_response(minted, principal, self._signing, issued_token_type=ACCESS_TOKEN_TOKEN_TYPE) + async def redeem_assertion( + self, assertion: str, client_id: str, client_secret: str | None, redeem_assertion: RedeemAssertion + ) -> Response: + """The RFC 7523 tail: prove the assertion and the requesting app, claim the assertion's + ``jti`` before anything is minted (RFC 7523 section 3 makes a replayed ``jti`` a + rejection, and the identity provider mints a fresh assertion cheaply, so a mint that + fails after the claim costs the client one more round trip, never a replay window), + then mint the proxy credential the user is revalidated for. No refresh token comes + with it (RFC 7521 section 4.1): the app renews by asking its identity provider for a + fresh assertion, which keeps the user's standing there in the loop.""" + target_refusal: Final = self._assertion_target_refusal() + if target_refusal is not None: + return target_refusal + identity: Final = await redeem_assertion(assertion, client_id, client_secret, self._request) + if isinstance(identity, SubjectTokenRefusal): + return _subject_token_refusal_response(identity) + refused: Final = _claim_refusal( + await self._guard.claim(identity.replay_key, identity.replay_ttl_seconds), + _oauth_error(400, "invalid_grant", "the assertion has already been redeemed"), + ) + if refused is not None: + return refused + minted: Final = await self._mint_proxy_credential( + identity.user_id, identity.team_id, allowed_routes=identity.allowed_routes + ) + if not isinstance(minted, MintedProxyCredential): + return _mint_failure_response(minted) + return assertion_credential_response(minted, identity.scope) + + def _assertion_target_refusal(self) -> Response | None: + """An ID-JAG names the proxy origin or its MCP resource (``{base}/mcp``) as ``aud``, and an RFC + 8707 client sends that same value as ``resource``, so both are this grant's targets.""" + if self._resource is None or is_proxy_api_resource(self._request, self._resource): + return None + if self._resource.rstrip("/") == f"{get_request_base_url(self._request)}/mcp": + return None + return _oauth_error( + 400, + "invalid_target", + "resource does not match the proxy API or MCP resource this gateway issues credentials for", + ) + def _proxy_api_target_refusal(self) -> Response | None: if self._resource is None or is_proxy_api_resource(self._request, self._resource): return None @@ -1807,6 +1928,19 @@ async def refresh_proxy_credential( ) +async def _jwt_bearer_grant( + request: Request, + assertion: str | None, + client_id: str, + client_secret: str | None, + redeem_assertion: RedeemAssertion, + issue: _GrantIssuer, +) -> Response: + if not assertion: + return _oauth_error(400, "invalid_request", "assertion is required") + return await issue.redeem_assertion(assertion, client_id, client_secret, redeem_assertion) + + async def _token_exchange_grant( subject_token: str | None, subject_token_type: str | None, diff --git a/litellm/proxy/_experimental/mcp_server/id_jag_grant.py b/litellm/proxy/_experimental/mcp_server/id_jag_grant.py new file mode 100644 index 00000000000..155d5280f5a --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/id_jag_grant.py @@ -0,0 +1,540 @@ +"""The RFC 7523 jwt-bearer grant on ``POST /token``: a requesting app that obtained an +ID-JAG (Identity Assertion Authorization Grant, the cross-app access profile) from the +customer's identity provider redeems it for the same proxy-API credential ``lite login`` +stores. The assertion is proven with the JWT auth issuer and JWKS plumbing the proxy +already trusts, its claims are checked the way the ID-JAG draft prescribes, the requesting +app must be one the operator allowlisted, and the ``jti`` is single-use until the +assertion expires.""" + +from __future__ import annotations + +import base64 +import binascii +import hashlib +import hmac +import math +from collections.abc import Awaitable, Iterator, Mapping, Sequence +from dataclasses import dataclass +from datetime import datetime, timezone +from functools import partial +from types import MappingProxyType +from typing import Final, Protocol +from urllib.parse import unquote_plus + +import jwt +from fastapi import Request + +from litellm._logging import verbose_proxy_logger +from litellm.models.user import LiteLLM_UserTable +from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( + AssertionIdentity, + ConsentTeam, + SubjectIdentity, + SubjectTokenRefusal, +) +from litellm.proxy._experimental.mcp_server.idp_token_exchange import ( + gateway_outage_refusal, + log_line, + read_token_exchange_prerequisites, +) +from litellm.proxy._experimental.mcp_server.oauth_utils import proxy_base_url_from_env +from litellm.proxy._types import IdJagRequestingApp +from litellm.proxy.auth.handle_jwt import ID_JAG_TYP, AssertionRejection, JWTHandler + +ID_JAG_REQUIRED_CLAIMS: Final = ("iss", "sub", "aud", "exp", "iat", "jti", "client_id") +REJECTED_ASSERTION: Final = "assertion was rejected" +TEAM_HEADER: Final = "x-litellm-team-id" +_INACTIVE_SUBJECT: Final = frozenset(("no_active_key", "unresolvable")) +_REPLAY_KEY_PREFIX: Final = "mcp_gateway_dcr:id_jag:" +SCOPE_MCP: Final = "litellm:mcp" +SCOPE_LLM: Final = "litellm:llm" +SCOPE_ROUTE_GROUPS: Final = MappingProxyType( + {SCOPE_MCP: ("mcp_inference_routes",), SCOPE_LLM: ("llm_api_routes_without_mcp",)} +) + + +@dataclass(frozen=True, slots=True) +class GrantedScope: + """What an assertion's ``scope`` claim buys: the scope echoed in the token response, exactly the + tokens enforced, and the route groups the credential is narrowed to. No claim means the full + credential, so both are ``None``.""" + + scope: str | None + allowed_routes: tuple[str, ...] | None + + +def granted_scope(claim: object, *, mcp_resource: bool) -> GrantedScope | SubjectTokenRefusal: + """The routes the assertion buys. With ``aud`` the bare origin the ``scope`` claim decides and + no claim is the full credential. With ``aud`` the MCP resource (``{base}/mcp``) the identity + provider named that resource, so the credential is ceilinged to the MCP routes whatever the + claim says: ``litellm:llm`` cannot be granted and is refused.""" + if claim is not None and not isinstance(claim, str): + return SubjectTokenRefusal(error="invalid_scope", description=_SCOPE_DESCRIPTION) + requested: Final = frozenset(claim.split()) if claim else frozenset() + if not requested <= SCOPE_ROUTE_GROUPS.keys(): + return SubjectTokenRefusal(error="invalid_scope", description=_SCOPE_DESCRIPTION) + if mcp_resource: + if SCOPE_LLM in requested: + return SubjectTokenRefusal(error="invalid_scope", description=_MCP_RESOURCE_DESCRIPTION) + return GrantedScope(scope=SCOPE_MCP, allowed_routes=SCOPE_ROUTE_GROUPS[SCOPE_MCP]) + if not requested: + return GrantedScope(scope=None, allowed_routes=None) + granted: Final = tuple(scope for scope in SCOPE_ROUTE_GROUPS if scope in requested) + return GrantedScope(scope=" ".join(granted), allowed_routes=tuple(_route_groups(granted))) + + +def names_mcp_resource(aud: object, base_url: str | None) -> bool: + """Whether the verified ``aud`` names only the MCP resource, ``{base}/mcp``, and not the + authorization server itself (the bare origin).""" + audiences: Final = frozenset((aud,) if isinstance(aud, str) else aud if isinstance(aud, list) else ()) + return base_url is not None and f"{base_url}/mcp" in audiences and base_url not in audiences + + +def _route_groups(scopes: Sequence[str]) -> Iterator[str]: + for scope in scopes: + yield from SCOPE_ROUTE_GROUPS[scope] + + +_SCOPE_DESCRIPTION: Final = ( + f"scope must be {SCOPE_MCP} (MCP gateway routes), {SCOPE_LLM} (LLM API routes), or both, space separated" +) +_MCP_RESOURCE_DESCRIPTION: Final = ( + f"the /mcp audience only grants MCP routes ({SCOPE_MCP}); use the bare origin as audience for LLM routes" +) + + +@dataclass(frozen=True, slots=True) +class IdJagPrerequisites: + """The deployment-level gates for redeeming an ID-JAG: the JWT auth whose issuer plumbing + proves the assertion must be on, backed by a database, and licensed, not mapping IdP tokens to + virtual keys (whose blocks and budgets a minted credential would sidestep), with a shared Redis + so an assertion is redeemed once across every worker and replica (the in-memory marker is per + process, so a second worker would mint again), ``PROXY_BASE_URL`` + must fix the gateway URL the assertion's ``aud`` has to name (a request's ``Host`` is the + caller's to choose), and the operator must have allowlisted the requesting apps. Discovery advertises the grant only when every gate + holds, and a redemption is refused naming the first one that does not.""" + + jwt_auth_enabled: bool + has_database: bool + licensed: bool + maps_jwts_to_virtual_keys: bool + shared_replay_store: bool + base_url: str | None + requesting_apps: tuple[IdJagRequestingApp, ...] + + @property + def available(self) -> bool: + return ( + self.jwt_auth_enabled + and self.has_database + and self.licensed + and not self.maps_jwts_to_virtual_keys + and self.shared_replay_store + and self.base_url is not None + and bool(self.requesting_apps) + ) + + def refusal(self) -> SubjectTokenRefusal | None: + if not self.jwt_auth_enabled: + return SubjectTokenRefusal( + error="unsupported_grant_type", + description="JWT auth is not enabled on this gateway, so it cannot redeem identity assertions", + ) + if not self.has_database: + return SubjectTokenRefusal( + error="unsupported_grant_type", + description="this gateway has no database, so it cannot redeem identity assertions", + ) + if not self.licensed: + return SubjectTokenRefusal( + error="unsupported_grant_type", + description="JWT auth is an enterprise only feature; no license is set", + ) + if self.maps_jwts_to_virtual_keys: + return SubjectTokenRefusal( + error="unsupported_grant_type", + description="this gateway maps IdP tokens to virtual keys, which the jwt-bearer grant does not serve", + ) + if not self.shared_replay_store: + return SubjectTokenRefusal( + error="unsupported_grant_type", + description="configure Redis (cache_params or coordination_redis) so this gateway can redeem an identity assertion exactly once across its workers", + ) + if self.base_url is None: + return SubjectTokenRefusal( + error="unsupported_grant_type", + description="set PROXY_BASE_URL to this gateway's public URL, the audience an identity assertion must name", + ) + if not self.requesting_apps: + return SubjectTokenRefusal( + error="unsupported_grant_type", + description="this gateway has no requesting apps allowlisted for identity assertions", + ) + return None + + +def read_id_jag_prerequisites() -> IdJagPrerequisites: + from litellm.proxy.proxy_server import ( # noqa: PLC0415 # rebound after startup, so read them per call + jwt_handler, + redis_usage_cache, + ) + + shared: Final = read_token_exchange_prerequisites() + return IdJagPrerequisites( + jwt_auth_enabled=shared.jwt_auth_enabled, + has_database=shared.has_database, + licensed=shared.licensed, + maps_jwts_to_virtual_keys=shared.maps_jwts_to_virtual_keys, + shared_replay_store=redis_usage_cache is not None, + base_url=proxy_base_url_from_env(), + requesting_apps=_requesting_apps(jwt_handler), + ) + + +def _requesting_apps(jwt_handler: JWTHandler) -> tuple[IdJagRequestingApp, ...]: + if not hasattr(jwt_handler, "litellm_jwtauth"): + return () + return tuple(jwt_handler.litellm_jwtauth.id_jag_requesting_apps or ()) + + +def jwt_bearer_available() -> bool: + return read_id_jag_prerequisites().available + + +def gateway_audiences(base_url: str | None) -> tuple[str, ...]: + """The issuer values this gateway answers to, which an ID-JAG's ``aud`` must name: its MCP + resource and its authorization server, both under the fixed ``PROXY_BASE_URL``.""" + return () if base_url is None else (f"{base_url}/mcp", base_url) + + +class DecodeAssertion(Protocol): + """Injected signature and claim verification ``(assertion, audience, required_claims)``: + ``JWTHandler.decode_trusted_assertion`` in production, which returns the verified claims, + a rejection, or raises when the gateway could not reach the issuer's keys.""" + + def __call__( + self, assertion: str, audience: Sequence[str], required_claims: Sequence[str], / + ) -> Awaitable[Mapping[str, object] | AssertionRejection]: ... + + +class ClaimReader(Protocol): + """The JWT auth user mapping as ``JWTHandler`` exposes it: the litellm user id and email the + verified claims name through ``user_id_jwt_field`` and ``user_email_jwt_field``.""" + + def get_user_id(self, token: dict[str, object], default_value: str | None) -> str | None: ... + + def get_user_email(self, token: dict[str, object], default_value: str | None) -> str | None: ... + + def get_all_jwt_team_ids(self, token: dict[str, object]) -> list[str]: ... + + def is_allowed_domain(self, user_email: str) -> bool: ... + + +class ResolveSubject(Protocol): + """Injected subject lookup ``(claims, team_header)``: the litellm user the verified claims + name through the JWT auth user mapping (``user_id_jwt_field`` or ``user_email_jwt_field``, + else ``sub``) and the team the credential is minted for (see ``select_team``), ``None`` when + the gateway knows no such user, or a refusal when the team cannot be settled. Raises on a + database outage.""" + + def __call__( + self, claims: Mapping[str, object], team_header: str | None, / + ) -> Awaitable[SubjectIdentity | SubjectTokenRefusal | None]: ... + + +class LookupUser(Protocol): + """Injected ``(user_id, user_email)`` lookup: the stored id of the existing litellm user the + claims name, ``None`` when the gateway knows no such user. Raises on a database outage.""" + + def __call__(self, user_id: str, user_email: str | None, /) -> Awaitable[str | None]: ... + + +class LookupTeams(Protocol): + """Injected lookup of the teams a user is on, or the failure text when they cannot be read.""" + + def __call__(self, user_id: str, /) -> Awaitable[Sequence[ConsentTeam] | str]: ... + + +class GetUser(Protocol): + def __call__( + self, *, user_id: str, user_email: str | None, sso_user_id: str + ) -> Awaitable[LiteLLM_UserTable | None]: ... + + +def proxy_user_lookup(get_user: GetUser) -> LookupUser: + """``LookupUser`` over the proxy's user table. The assertion's subject is matched as the + local ``user_id`` and as the ``sso_user_id`` an SSO login linked, the way the JWT and SSO + login paths match it, so a linked account needs no email claim.""" + + async def lookup_user(user_id: str, user_email: str | None) -> str | None: + user: Final = await get_user(user_id=user_id, user_email=user_email, sso_user_id=user_id) + return None if user is None else user.user_id + + return lookup_user + + +def subject_resolver(reader: ClaimReader, lookup_user: LookupUser, lookup_teams: LookupTeams) -> ResolveSubject: + """``ResolveSubject`` over the JWT auth claim mapping and the proxy's user and team lookups.""" + + async def resolve_subject( + claims: Mapping[str, object], team_header: str | None + ) -> SubjectIdentity | SubjectTokenRefusal | None: + token: Final = dict(claims) + subject: Final = claims.get("sub") + claimed_user_id: Final = reader.get_user_id(token=token, default_value=None) or ( + subject if isinstance(subject, str) else None + ) + if claimed_user_id is None: + return None + user_email: Final = reader.get_user_email(token=token, default_value=None) + if not reader.is_allowed_domain(user_email or ""): + return SubjectTokenRefusal( + error="invalid_grant", + description="the assertion's subject email is missing or outside the allowed domain", + ) + user_id: Final = await lookup_user(claimed_user_id, user_email) + if user_id is None: + return None + teams: Final = await lookup_teams(user_id) + if teams in _INACTIVE_SUBJECT: + return SubjectTokenRefusal( + error="invalid_grant", description="the assertion's subject is not an active user" + ) + if isinstance(teams, str): + return SubjectTokenRefusal( + error="temporarily_unavailable", description="the gateway could not read this user's teams; retry" + ) + team: Final = select_team(team_header, tuple(reader.get_all_jwt_team_ids(token=token)), teams) + if isinstance(team, SubjectTokenRefusal): + return team + return SubjectIdentity(user_id=user_id, team_id=team) + + return resolve_subject + + +def select_team( + team_header: str | None, claim_teams: Sequence[str], teams: Sequence[ConsentTeam] +) -> str | None | SubjectTokenRefusal: + """The team an ID-JAG redemption mints for, settled before the assertion is consumed since + the minter refuses a teamless credential for a user who has a team. The ``x-litellm-team-id`` + header (a team id or alias) wins and must name one of the user's teams, then the assertion's + mapped team claim, then the user's only team. A user on several teams must name one.""" + if team_header is not None: + named: Final = next((team for team in teams if team_header in (team.team_id, team.team_alias)), None) + if named is None: + return SubjectTokenRefusal( + error="invalid_request", description=f"{TEAM_HEADER} does not name one of this user's teams" + ) + return named.team_id + asserted: Final = tuple(team.team_id for team in teams if team.team_id in claim_teams) + if len(asserted) == 1: + return asserted[0] + if len(asserted) > 1: + return SubjectTokenRefusal( + error="invalid_request", + description=f"the assertion names several of this user's teams; send the {TEAM_HEADER} header naming the team for this credential", + ) + if len(teams) == 1: + return teams[0].team_id + if teams: + return SubjectTokenRefusal( + error="invalid_request", + description=f"this user belongs to several teams; send the {TEAM_HEADER} header naming the team for this credential", + ) + return None + + +@dataclass(frozen=True, slots=True) +class AssertionRedemption: + """One ``jwt-bearer`` request as the grant sees it: the assertion, how the requesting app + identified and authenticated itself, the team the credential is asked for, the issuer + values this gateway answers to (what the assertion's ``aud`` must name), the clock, and the + leeway the verifier grants ``exp``, so the replay marker outlives every moment the assertion + would still verify.""" + + assertion: str + client_id: str + client_secret: str | None + authorization_header: str | None + team_id: str | None + audiences: tuple[str, ...] + now: float + leeway_seconds: float + + +async def redeem_id_jag_assertion( + assertion: str, client_id: str, client_secret: str | None, request: Request +) -> AssertionIdentity | SubjectTokenRefusal: + from litellm.proxy._experimental.mcp_server.proxy_api_credentials import ( # noqa: PLC0415 # circular import at module load + lookup_consent_teams_from_database, + ) + from litellm.proxy.auth.auth_checks import get_user_object # noqa: PLC0415 # circular import at module load + from litellm.proxy.proxy_server import ( # noqa: PLC0415 # rebound after startup, so read them per call + jwt_handler, + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + resolve_subject: Final = subject_resolver( + _claim_reader(jwt_handler), + proxy_user_lookup( + partial( + get_user_object, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + proxy_logging_obj=proxy_logging_obj, + check_db_only=True, + ) + ), + lookup_consent_teams_from_database, + ) + prerequisites: Final = read_id_jag_prerequisites() + return await identity_from_id_jag( + AssertionRedemption( + assertion=assertion, + client_id=client_id, + client_secret=client_secret, + authorization_header=request.headers.get("authorization"), + team_id=request.headers.get(TEAM_HEADER), + audiences=gateway_audiences(prerequisites.base_url), + now=datetime.now(timezone.utc).timestamp(), + leeway_seconds=jwt_handler.leeway, + ), + prerequisites=prerequisites, + decode=jwt_handler.decode_trusted_assertion, + resolve_subject=resolve_subject, + ) + + +def _claim_reader(handler: ClaimReader) -> ClaimReader: + return handler + + +async def identity_from_id_jag( + redemption: AssertionRedemption, + prerequisites: IdJagPrerequisites, + decode: DecodeAssertion, + resolve_subject: ResolveSubject, +) -> AssertionIdentity | SubjectTokenRefusal: + """Prove an ID-JAG for the grant tail. The requesting app is checked first, the way a token + endpoint authenticates its client before it looks at the grant, so an app the operator never + allowlisted learns nothing about the assertion. Everything wrong with the assertion itself is + ``invalid_grant`` (RFC 7523 section 3.1), a check the gateway could not complete is + ``temporarily_unavailable``, and the reason stays in the proxy log because this endpoint is + public and the verifier's wording can quote the assertion.""" + unmet: Final = prerequisites.refusal() + if unmet is not None: + return unmet + app: Final = next((app for app in prerequisites.requesting_apps if app.client_id == redemption.client_id), None) + if app is None: + return SubjectTokenRefusal(error="invalid_client", description="client_id is not an allowlisted requesting app") + if not _client_authenticated(app, redemption.client_secret, redemption.authorization_header): + return SubjectTokenRefusal(error="invalid_client", description="client authentication failed") + try: + header: Final = jwt.get_unverified_header(redemption.assertion) + except jwt.PyJWTError: + return SubjectTokenRefusal(error="invalid_grant", description="assertion is not a JWT") + if header.get("typ") != ID_JAG_TYP: + return SubjectTokenRefusal(error="invalid_grant", description=f"assertion typ must be {ID_JAG_TYP}") + try: + verified: Final = await decode(redemption.assertion, redemption.audiences, ID_JAG_REQUIRED_CLAIMS) + except Exception as denied: # noqa: BLE001 # the verifier raises its 503 for a JWKS outage, and anything else is a bad assertion + return gateway_outage_refusal(denied, denied, what="an assertion", parameter="assertion") or _rejected( + str(denied) + ) + if isinstance(verified, AssertionRejection): + return _rejected(verified.reason) + marker: Final = _replay_marker(verified, redemption) + if isinstance(marker, SubjectTokenRefusal): + return marker + replay_key, replay_ttl_seconds = marker + scope: Final = granted_scope( + verified.get("scope"), mcp_resource=names_mcp_resource(verified.get("aud"), prerequisites.base_url) + ) + if isinstance(scope, SubjectTokenRefusal): + return scope + try: + subject: Final = await resolve_subject(verified, redemption.team_id) + except Exception as failed: # noqa: BLE001 # the user lookup raises for a missing user and for a database outage alike + return gateway_outage_refusal( + failed, failed, what="an assertion's subject", parameter="assertion" + ) or _rejected(f"subject lookup failed: {failed}") + if subject is None: + return _rejected("the assertion names no user the gateway knows") + if isinstance(subject, SubjectTokenRefusal): + return subject + return AssertionIdentity( + user_id=subject.user_id, + team_id=subject.team_id, + replay_key=replay_key, + replay_ttl_seconds=replay_ttl_seconds, + scope=scope.scope, + allowed_routes=scope.allowed_routes, + ) + + +def _replay_marker( + verified: Mapping[str, object], redemption: AssertionRedemption +) -> tuple[str, int] | SubjectTokenRefusal: + """The single-use key and its lifetime for a verified assertion, once the assertion is bound to + the presenting app, carries no proof-of-possession the gateway cannot check, and is still alive.""" + if verified.get("client_id") != redemption.client_id: + return _rejected("the assertion's client_id is not the requesting app that presented it") + if "cnf" in verified: + return _rejected("the assertion is proof-bound and this gateway does not check proof of possession") + issuer: Final = verified.get("iss") + jti: Final = verified.get("jti") + exp: Final = verified.get("exp") + if not isinstance(issuer, str) or not isinstance(jti, str) or not jti or not isinstance(exp, (int, float)): + return _rejected("the assertion's iss, jti, or exp claim has the wrong shape") + replay_ttl_seconds: Final = math.ceil(exp + redemption.leeway_seconds - redemption.now) + if replay_ttl_seconds <= 0: + return _rejected("the assertion has expired") + return _REPLAY_KEY_PREFIX + hashlib.sha256(f"{issuer}\n{jti}".encode()).hexdigest(), replay_ttl_seconds + + +def _rejected(reason: str) -> SubjectTokenRefusal: + verbose_proxy_logger.warning("jwt-bearer grant refused an assertion: %s", log_line(reason)) + return SubjectTokenRefusal(error="invalid_grant", description=REJECTED_ASSERTION) + + +def _client_authenticated(app: IdJagRequestingApp, form_secret: str | None, authorization_header: str | None) -> bool: + """A public app (no secret configured) is identified by ``client_id`` alone. A confidential + app must present its secret, by HTTP Basic (RFC 6749 section 2.3.1, the credentials + form-urlencoded inside the base64) or the ``client_secret`` form field, compared in constant + time.""" + if app.client_secret is None: + return True + basic: Final = _basic_credentials(authorization_header) + if basic is not None: + return basic[0] == app.client_id and _same_secret(basic[1], app.client_secret) + return form_secret is not None and _same_secret(form_secret, app.client_secret) + + +def _same_secret(presented: str, expected: str) -> bool: + return hmac.compare_digest(presented.encode(), expected.encode()) + + +def basic_client_id(authorization_header: str | None) -> str | None: + """The ``client_id`` an HTTP Basic ``Authorization`` header carries, for a confidential app + that leaves the form field out as RFC 6749 section 2.3.1 allows.""" + credentials: Final = _basic_credentials(authorization_header) + return None if credentials is None else credentials[0] + + +def _basic_credentials(authorization_header: str | None) -> tuple[str, str] | None: + if authorization_header is None: + return None + scheme, _, encoded = authorization_header.strip().partition(" ") + if scheme.lower() != "basic" or not encoded: + return None + try: + decoded: Final = base64.b64decode(encoded.strip(), validate=True).decode() + except (binascii.Error, UnicodeDecodeError): + return None + client_id, separator, secret = decoded.partition(":") + if not separator: + return None + return unquote_plus(client_id), unquote_plus(secret) diff --git a/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py b/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py index 15632eb4783..5044a5bd3dc 100644 --- a/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py +++ b/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py @@ -20,13 +20,15 @@ from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler EXCHANGE_ROUTE: Final = "/token" REJECTED_SUBJECT_TOKEN: Final = "subject_token was rejected by the gateway's JWT auth" -SUBJECT_TOKEN_CHECK_UNAVAILABLE: Final = ( - "the gateway could not verify subject_token because its identity provider or database is unavailable; retry" +_CHECK_UNAVAILABLE: Final = ( + "the gateway could not verify {parameter} because its identity provider or database is unavailable; retry" ) -SUBJECT_TOKEN_CHECK_FAULTED: Final = ( - "the gateway could not verify subject_token because its database reported a fault that is not a transient " +_CHECK_FAULTED: Final = ( + "the gateway could not verify {parameter} because its database reported a fault that is not a transient " "outage; retrying will not help until the gateway deployment is repaired" ) +SUBJECT_TOKEN_CHECK_UNAVAILABLE: Final = _CHECK_UNAVAILABLE.format(parameter="subject_token") +SUBJECT_TOKEN_CHECK_FAULTED: Final = _CHECK_FAULTED.format(parameter="subject_token") GatewayOutage = Literal["retryable", "faulted"] @@ -182,20 +184,37 @@ async def identity_from_subject_token( def _refusal_for(denied: Exception, reason: object) -> SubjectTokenRefusal: + outage: Final = gateway_outage_refusal(denied, reason, what="a subject_token", parameter="subject_token") + if outage is not None: + return outage + verbose_proxy_logger.warning("token exchange refused a subject_token: %s", log_line(reason)) + return SubjectTokenRefusal(error="invalid_request", description=REJECTED_SUBJECT_TOKEN) + + +def gateway_outage_refusal(denied: Exception, reason: object, what: str, parameter: str) -> SubjectTokenRefusal | None: + """The ``temporarily_unavailable`` refusal when ``denied`` is the gateway failing to check + ``what`` (its database or the IdP's JWKS down), naming the grant ``parameter`` it could not verify, + else ``None`` because the token itself was bad.""" outage: Final = _gateway_could_not_verify(denied) if outage is None: - verbose_proxy_logger.warning("token exchange refused a subject_token: %s", reason) - return SubjectTokenRefusal(error="invalid_request", description=REJECTED_SUBJECT_TOKEN) - verbose_proxy_logger.error("token exchange could not verify a subject_token, %s: %s", outage, reason) - return SubjectTokenRefusal(error="temporarily_unavailable", description=_check_unavailable_description(outage)) + return None + verbose_proxy_logger.error("token endpoint could not verify %s, %s: %s", what, outage, log_line(reason)) + return SubjectTokenRefusal( + error="temporarily_unavailable", description=_check_unavailable_description(outage, parameter) + ) -def _check_unavailable_description(outage: GatewayOutage) -> str: +def log_line(reason: object) -> str: + """``reason`` as one log line, so text a caller controls cannot forge log entries.""" + return str(reason).replace("\r", " ").replace("\n", " ") + + +def _check_unavailable_description(outage: GatewayOutage, parameter: str) -> str: match outage: case "retryable": - return SUBJECT_TOKEN_CHECK_UNAVAILABLE + return _CHECK_UNAVAILABLE.format(parameter=parameter) case "faulted": - return SUBJECT_TOKEN_CHECK_FAULTED + return _CHECK_FAULTED.format(parameter=parameter) case _: assert_never(outage) diff --git a/litellm/proxy/_experimental/mcp_server/oauth_utils.py b/litellm/proxy/_experimental/mcp_server/oauth_utils.py index b0448a62c9e..69732186d75 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_utils.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_utils.py @@ -128,6 +128,11 @@ def build_upstream_authorize_url(authorization_url: str, params: Mapping[str, st return urlunparse(parsed._replace(query=urlencode(merged))) +def proxy_base_url_from_env() -> str | None: + """The gateway's fixed public URL from ``PROXY_BASE_URL``, ``None`` when unset or malformed.""" + return _resolve_proxy_base_url_env() + + def _resolve_proxy_base_url_env() -> str | None: global _warned_invalid_proxy_base_url configured: Final = os.environ.get("PROXY_BASE_URL", "").strip() diff --git a/litellm/proxy/_experimental/mcp_server/proxy_api_credentials.py b/litellm/proxy/_experimental/mcp_server/proxy_api_credentials.py index 2f7fcaef645..fd637a2f187 100644 --- a/litellm/proxy/_experimental/mcp_server/proxy_api_credentials.py +++ b/litellm/proxy/_experimental/mcp_server/proxy_api_credentials.py @@ -25,7 +25,16 @@ from litellm.proxy.management_endpoints.ui_sso import ( async def lookup_consent_teams(user_id: str) -> tuple[ConsentTeam, ...] | ReloadUserFailure: - user: Final = await load_active_user_by_id(user_id) + return await _consent_teams(await load_active_user_by_id(user_id)) + + +async def lookup_consent_teams_from_database(user_id: str) -> tuple[ConsentTeam, ...] | ReloadUserFailure: + """The same teams read from the database row, for a grant that is consumed once it is settled: the + membership the minter will re-check must be the membership the team was picked from.""" + return await _consent_teams(await load_active_user_by_id(user_id, source="database")) + + +async def _consent_teams(user: LiteLLM_UserTable | ReloadUserFailure) -> tuple[ConsentTeam, ...] | ReloadUserFailure: if isinstance(user, str): return user details: Final = await _team_details(user.teams) @@ -39,7 +48,7 @@ async def lookup_consent_teams(user_id: str) -> tuple[ConsentTeam, ...] | Reload async def mint_proxy_credential( - user_id: str, team_id: str | None + user_id: str, team_id: str | None, *, allowed_routes: Sequence[str] | None = None ) -> MintedProxyCredential | ProxyCredentialMintFailure: """Mint the ``lite login`` credential for a consented grant. Membership is checked live against the database row, so a team the user left between consent and redemption (or between refreshes) @@ -53,7 +62,9 @@ async def mint_proxy_credential( user row handed to the minter carries no team list, exactly like ``lite login``'s, so the minter's own first-team fallback stays inert. The credential carries the role the proxy already enforces for the user on every request, so a row with no role (JWT auth's - upsert writes none) mints as an internal user instead of being refused.""" + upsert writes none) mints as an internal user instead of being refused. ``allowed_routes`` + narrows the credential to those route groups, the way a virtual key's ``allowed_routes`` does, + for a grant whose scope asked for less than the full credential.""" user: Final = await load_active_user_by_id(user_id, source="database") if isinstance(user, str): return user @@ -75,6 +86,7 @@ async def mint_proxy_credential( team_alias=selected.team_alias, team_models=selected.team_models, team_model_aliases=selected.team_model_aliases, + allowed_routes=allowed_routes, ) return MintedProxyCredential( key=key, diff --git a/litellm/proxy/_lazy_features.py b/litellm/proxy/_lazy_features.py index a909a7a993a..a7adb5c1bf5 100644 --- a/litellm/proxy/_lazy_features.py +++ b/litellm/proxy/_lazy_features.py @@ -155,18 +155,13 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = ( module_path="litellm.proxy.management_endpoints.mcp_management_endpoints", path_prefixes=("/v1/mcp/",), ), - LazyFeature( - # Also serves /.well-known/oauth-* (OAuth metadata discovery). - # No /mcp/oauth prefix here: the mounted /mcp sub-app would - # shadow it, and there are no actual routes there anyway. - name="mcp_byok_oauth", - module_path="litellm.proxy._experimental.mcp_server.byok_oauth_endpoints", - path_prefixes=("/v1/mcp/oauth", "/.well-known/oauth-"), - ), LazyFeature( # Serves OAuth dance endpoints (/authorize, /token, /callback, # /register) plus several /.well-known/ discovery URLs at the proxy # root — needed for MCP-over-OAuth flows even before /mcp is hit. + # Registered before mcp_byok_oauth, which also claims the root + # /.well-known/oauth-authorization-server: the root document is the + # gateway's own authorization server metadata (token endpoint /token). name="mcp_discoverable", module_path="litellm.proxy._experimental.mcp_server.discoverable_endpoints", path_prefixes=( @@ -185,6 +180,14 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = ( # Catches the /{mcp_server_name}/authorize|token|register variants. path_suffixes=("/authorize", "/token", "/register"), ), + LazyFeature( + # Also serves /.well-known/oauth-* (OAuth metadata discovery). + # No /mcp/oauth prefix here: the mounted /mcp sub-app would + # shadow it, and there are no actual routes there anyway. + name="mcp_byok_oauth", + module_path="litellm.proxy._experimental.mcp_server.byok_oauth_endpoints", + path_prefixes=("/v1/mcp/oauth", "/.well-known/oauth-"), + ), LazyFeature( name="mcp_rest", module_path="litellm.proxy._experimental.mcp_server.rest_endpoints", diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 20819a59a43..2eb62a07447 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -35771,7 +35771,7 @@ "/.well-known/oauth-authorization-server": { "get": { "description": "OAuth authorization server discovery endpoint.\n\nSupports both legacy pattern (/{server_name}) and root endpoint.", - "operationId": "oauth_authorization_server_mcp__well_known_oauth_authorization_server_get", + "operationId": "oauth_authorization_server_mcp__well_known_oauth_authorization_server_get_2", "parameters": [ { "in": "query", @@ -35819,7 +35819,7 @@ "/.well-known/oauth-authorization-server/mcp": { "get": { "description": "OAuth authorization server discovery for the aggregate /mcp endpoint, the RFC 8414\npath-inserted form for a client that treats {base}/mcp as its authorization base URL.\n\nThe single-segment /mcp is reserved for the aggregate so the discovery chain stays\nconsistent: the aggregate protected-resource document advertises {base}/mcp as its\nauthorization server, so the document served here must have issuer {base}/mcp. A server\nliterally named ``mcp`` therefore does not take this route; it keeps its standard\ntwo-segment discovery at /.well-known/oauth-authorization-server/mcp/mcp. Letting the\nper-server row win here instead would serve an issuer of {base} against a resource that\nadvertised {base}/mcp, which fails the RFC 8414 issuer check and breaks the front door.", - "operationId": "oauth_authorization_server_aggregate__well_known_oauth_authorization_server_mcp_get", + "operationId": "oauth_authorization_server_aggregate__well_known_oauth_authorization_server_mcp_get_2", "responses": { "200": { "content": { @@ -35839,7 +35839,7 @@ "/.well-known/oauth-authorization-server/mcp/{mcp_server_name}": { "get": { "description": "OAuth authorization server discovery endpoint using standard MCP URL pattern.\n\nStandard pattern: /mcp/{server_name}\nDiscovery path: /.well-known/oauth-authorization-server/mcp/{server_name}", - "operationId": "oauth_authorization_server_mcp_standard__well_known_oauth_authorization_server_mcp__mcp_server_name__get", + "operationId": "oauth_authorization_server_mcp_standard__well_known_oauth_authorization_server_mcp__mcp_server_name__get_2", "parameters": [ { "in": "path", @@ -35880,7 +35880,7 @@ "/.well-known/oauth-authorization-server/{mcp_server_name}": { "get": { "description": "OAuth authorization server discovery endpoint.\n\nSupports both legacy pattern (/{server_name}) and root endpoint.", - "operationId": "oauth_authorization_server_mcp__well_known_oauth_authorization_server__mcp_server_name__get", + "operationId": "oauth_authorization_server_mcp__well_known_oauth_authorization_server__mcp_server_name__get_2", "parameters": [ { "in": "path", @@ -35928,7 +35928,7 @@ "/.well-known/oauth-authorization-server/{mcp_server_name}/mcp": { "get": { "description": "OAuth authorization server discovery for legacy /{server_name}/mcp pattern.", - "operationId": "oauth_authorization_server_legacy__well_known_oauth_authorization_server__mcp_server_name__mcp_get", + "operationId": "oauth_authorization_server_legacy__well_known_oauth_authorization_server__mcp_server_name__mcp_get_2", "parameters": [ { "in": "path", @@ -35968,7 +35968,7 @@ }, "/.well-known/oauth-protected-resource": { "get": { - "operationId": "oauth_protected_resource_root__well_known_oauth_protected_resource_get", + "operationId": "oauth_protected_resource_root__well_known_oauth_protected_resource_get_2", "responses": { "200": { "content": { @@ -36004,7 +36004,7 @@ "/.well-known/oauth-protected-resource/mcp": { "get": { "description": "OAuth protected resource discovery for the aggregate /mcp endpoint.\n\nThe single-segment ``/mcp`` path does not collide with any per-server PRM pattern\n(those are two-segment: ``/mcp/{server}`` or ``/{server}/mcp``), so this unambiguously\ndescribes the aggregate resource.", - "operationId": "oauth_protected_resource_aggregate__well_known_oauth_protected_resource_mcp_get", + "operationId": "oauth_protected_resource_aggregate__well_known_oauth_protected_resource_mcp_get_2", "responses": { "200": { "content": { @@ -36024,7 +36024,7 @@ "/.well-known/oauth-protected-resource/mcp/{mcp_server_name}": { "get": { "description": "OAuth protected resource discovery endpoint using standard MCP URL pattern.\n\nStandard pattern: /mcp/{server_name}\nDiscovery path: /.well-known/oauth-protected-resource/mcp/{server_name}\n\nThis endpoint is compliant with MCP specification and works with standard\nMCP clients like mcp-inspector and VSCode Copilot.", - "operationId": "oauth_protected_resource_mcp_standard__well_known_oauth_protected_resource_mcp__mcp_server_name__get", + "operationId": "oauth_protected_resource_mcp_standard__well_known_oauth_protected_resource_mcp__mcp_server_name__get_2", "parameters": [ { "in": "path", @@ -36065,7 +36065,7 @@ "/.well-known/oauth-protected-resource/{mcp_server_name}/mcp": { "get": { "description": "OAuth protected resource discovery endpoint using LiteLLM legacy URL pattern.\n\nLegacy pattern: /{server_name}/mcp\nDiscovery path: /.well-known/oauth-protected-resource/{server_name}/mcp\n\nThis endpoint is kept for backward compatibility. New integrations should\nuse the standard MCP pattern (/mcp/{server_name}) instead.", - "operationId": "oauth_protected_resource_mcp__well_known_oauth_protected_resource__mcp_server_name__mcp_get", + "operationId": "oauth_protected_resource_mcp__well_known_oauth_protected_resource__mcp_server_name__mcp_get_2", "parameters": [ { "in": "path", @@ -36194,9 +36194,27 @@ }, "Body_token_endpoint__mcp_server_name__token_post": { "properties": { + "assertion": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Assertion" + }, "client_id": { - "title": "Client Id", - "type": "string" + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Client Id" }, "client_secret": { "anyOf": [ @@ -36293,17 +36311,34 @@ } }, "required": [ - "grant_type", - "client_id" + "grant_type" ], "title": "Body_token_endpoint__mcp_server_name__token_post", "type": "object" }, "Body_token_endpoint_token_post": { "properties": { + "assertion": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Assertion" + }, "client_id": { - "title": "Client Id", - "type": "string" + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Client Id" }, "client_secret": { "anyOf": [ @@ -36400,8 +36435,7 @@ } }, "required": [ - "grant_type", - "client_id" + "grant_type" ], "title": "Body_token_endpoint_token_post", "type": "object" @@ -38151,7 +38185,7 @@ "/.well-known/oauth-authorization-server": { "get": { "description": "OAuth authorization server discovery endpoint.\n\nSupports both legacy pattern (/{server_name}) and root endpoint.", - "operationId": "oauth_authorization_server_mcp__well_known_oauth_authorization_server_get_2", + "operationId": "oauth_authorization_server_mcp__well_known_oauth_authorization_server_get", "parameters": [ { "in": "query", @@ -38199,7 +38233,7 @@ "/.well-known/oauth-authorization-server/mcp": { "get": { "description": "OAuth authorization server discovery for the aggregate /mcp endpoint, the RFC 8414\npath-inserted form for a client that treats {base}/mcp as its authorization base URL.\n\nThe single-segment /mcp is reserved for the aggregate so the discovery chain stays\nconsistent: the aggregate protected-resource document advertises {base}/mcp as its\nauthorization server, so the document served here must have issuer {base}/mcp. A server\nliterally named ``mcp`` therefore does not take this route; it keeps its standard\ntwo-segment discovery at /.well-known/oauth-authorization-server/mcp/mcp. Letting the\nper-server row win here instead would serve an issuer of {base} against a resource that\nadvertised {base}/mcp, which fails the RFC 8414 issuer check and breaks the front door.", - "operationId": "oauth_authorization_server_aggregate__well_known_oauth_authorization_server_mcp_get_2", + "operationId": "oauth_authorization_server_aggregate__well_known_oauth_authorization_server_mcp_get", "responses": { "200": { "content": { @@ -38219,7 +38253,7 @@ "/.well-known/oauth-authorization-server/mcp/{mcp_server_name}": { "get": { "description": "OAuth authorization server discovery endpoint using standard MCP URL pattern.\n\nStandard pattern: /mcp/{server_name}\nDiscovery path: /.well-known/oauth-authorization-server/mcp/{server_name}", - "operationId": "oauth_authorization_server_mcp_standard__well_known_oauth_authorization_server_mcp__mcp_server_name__get_2", + "operationId": "oauth_authorization_server_mcp_standard__well_known_oauth_authorization_server_mcp__mcp_server_name__get", "parameters": [ { "in": "path", @@ -38260,7 +38294,7 @@ "/.well-known/oauth-authorization-server/{mcp_server_name}": { "get": { "description": "OAuth authorization server discovery endpoint.\n\nSupports both legacy pattern (/{server_name}) and root endpoint.", - "operationId": "oauth_authorization_server_mcp__well_known_oauth_authorization_server__mcp_server_name__get_2", + "operationId": "oauth_authorization_server_mcp__well_known_oauth_authorization_server__mcp_server_name__get", "parameters": [ { "in": "path", @@ -38308,7 +38342,7 @@ "/.well-known/oauth-authorization-server/{mcp_server_name}/mcp": { "get": { "description": "OAuth authorization server discovery for legacy /{server_name}/mcp pattern.", - "operationId": "oauth_authorization_server_legacy__well_known_oauth_authorization_server__mcp_server_name__mcp_get_2", + "operationId": "oauth_authorization_server_legacy__well_known_oauth_authorization_server__mcp_server_name__mcp_get", "parameters": [ { "in": "path", @@ -38348,7 +38382,7 @@ }, "/.well-known/oauth-protected-resource": { "get": { - "operationId": "oauth_protected_resource_root__well_known_oauth_protected_resource_get_2", + "operationId": "oauth_protected_resource_root__well_known_oauth_protected_resource_get", "responses": { "200": { "content": { @@ -38384,7 +38418,7 @@ "/.well-known/oauth-protected-resource/mcp": { "get": { "description": "OAuth protected resource discovery for the aggregate /mcp endpoint.\n\nThe single-segment ``/mcp`` path does not collide with any per-server PRM pattern\n(those are two-segment: ``/mcp/{server}`` or ``/{server}/mcp``), so this unambiguously\ndescribes the aggregate resource.", - "operationId": "oauth_protected_resource_aggregate__well_known_oauth_protected_resource_mcp_get_2", + "operationId": "oauth_protected_resource_aggregate__well_known_oauth_protected_resource_mcp_get", "responses": { "200": { "content": { @@ -38404,7 +38438,7 @@ "/.well-known/oauth-protected-resource/mcp/{mcp_server_name}": { "get": { "description": "OAuth protected resource discovery endpoint using standard MCP URL pattern.\n\nStandard pattern: /mcp/{server_name}\nDiscovery path: /.well-known/oauth-protected-resource/mcp/{server_name}\n\nThis endpoint is compliant with MCP specification and works with standard\nMCP clients like mcp-inspector and VSCode Copilot.", - "operationId": "oauth_protected_resource_mcp_standard__well_known_oauth_protected_resource_mcp__mcp_server_name__get_2", + "operationId": "oauth_protected_resource_mcp_standard__well_known_oauth_protected_resource_mcp__mcp_server_name__get", "parameters": [ { "in": "path", @@ -38445,7 +38479,7 @@ "/.well-known/oauth-protected-resource/{mcp_server_name}/mcp": { "get": { "description": "OAuth protected resource discovery endpoint using LiteLLM legacy URL pattern.\n\nLegacy pattern: /{server_name}/mcp\nDiscovery path: /.well-known/oauth-protected-resource/{server_name}/mcp\n\nThis endpoint is kept for backward compatibility. New integrations should\nuse the standard MCP pattern (/mcp/{server_name}) instead.", - "operationId": "oauth_protected_resource_mcp__well_known_oauth_protected_resource__mcp_server_name__mcp_get_2", + "operationId": "oauth_protected_resource_mcp__well_known_oauth_protected_resource__mcp_server_name__mcp_get", "parameters": [ { "in": "path", diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d62c6b0d51f..2c98cfeb33d 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -551,6 +551,7 @@ class LiteLLMRoutes(enum.Enum): + list(agent_inference_routes) + model_info_routes ) + llm_api_routes_without_mcp = sorted(set(llm_api_routes) - set(mcp_inference_routes)) info_routes = [ "/key/info", "/key/health", @@ -4762,6 +4763,26 @@ class JWTIssuerConfig(LiteLLMBaseModel): DEFAULT_JWKS_STALE_TTL: Final = 3600 +class IdJagRequestingApp(LiteLLMBaseModel): + """A requesting application the operator lets redeem an ID-JAG (Identity Assertion + Authorization Grant) at the MCP gateway's ``POST /token``. ``client_id`` is the value + the identity provider puts in the assertion's ``client_id`` claim and the app sends as + the ``client_id`` form field. A ``client_secret`` makes the app a confidential client + that must authenticate with it (HTTP Basic or the ``client_secret`` form field); without + one the app is a public client.""" + + model_config = ConfigDict(frozen=True) + + client_id: str = Field( + min_length=1, description="The requesting app's client_id as the identity provider asserts it." + ) + client_secret: str | None = Field( + default=None, + min_length=1, + description="Shared secret the app must present at the gateway's token endpoint. Omit for a public client.", + ) + + class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): """ A class to define the roles and permissions for a LiteLLM Proxy w/ JWT Auth. @@ -4942,6 +4963,10 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): default=None, description="Optional issuer-bound JWT validation rules. When a token's `iss` matches a configured issuer, validation uses that issuer's JWKS, audience, and claim mappings. Tokens with an unlisted `iss` fall back to the global JWT_AUDIENCE/JWT_ISSUER validation path — this is additive routing, not an allow-list.", ) + id_jag_requesting_apps: tuple[IdJagRequestingApp, ...] | None = Field( + default=None, + description="Requesting apps allowed to redeem an ID-JAG at the MCP gateway's POST /token with the RFC 7523 jwt-bearer grant. Unset or empty turns that grant off.", + ) ######################################################### def __init__(self, **kwargs: Any) -> None: diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 8b32eb450e8..b395a826625 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -3960,6 +3960,7 @@ class ExperimentalUIJWTToken: team_models: Sequence[str] | None = None, team_model_aliases: Mapping[str, str] | None = None, max_budget: float | None = None, + allowed_routes: Sequence[str] | None = None, ) -> str: """ Generate a JWT token for CLI authentication with configurable expiration. @@ -3973,6 +3974,8 @@ class ExperimentalUIJWTToken: team_alias: Team alias for the selected team, if available team_models: Model allowlist granted by the selected team team_model_aliases: Team model aliases for the selected team + allowed_routes: Route groups the session may call, enforced like a virtual key's + ``allowed_routes`` on every route; None leaves every route the role allows Returns: Encrypted JWT token string @@ -4017,6 +4020,7 @@ class ExperimentalUIJWTToken: max_parallel_requests=None, user_role=LitellmUserRoles(user_info.user_role), is_session_token=True, + allowed_routes=list(allowed_routes) if allowed_routes is not None else None, ) return encrypt_bearer_token(valid_token.model_dump_json(exclude_none=True), prefix=LITELLM_SESSION_TOKEN_PREFIX) diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 2421eb10c41..05b8cca19af 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -35,6 +35,7 @@ from litellm.llms.custom_httpx.httpx_handler import HTTPHandler from litellm.proxy._types import ( DEFAULT_JWKS_STALE_TTL, RBAC_ROLES, + IdJagRequestingApp, JWKKeyValue, JWTAuthBuilderResult, JWTIssuerConfig, @@ -99,6 +100,10 @@ class NoMatchingJWTPublicKeyError(Exception): """Raised when a JWKS endpoint returns no key matching the requested ``kid``.""" +class IdJagIsNotABearerTokenError(Exception): + """Raised when bearer auth is handed an ID-JAG, which only the token endpoint may redeem.""" + + class JWKSUnreachableError(Exception): """Raised when an IdP's JWKS / OIDC discovery endpoint is unreachable and no cached copy is left to fall back on.""" @@ -114,9 +119,23 @@ _CACHED_VALUE: Final = TypeAdapter(object) _CachedValueT = TypeVar("_CachedValueT", bound=JWKKeyValue | str) +ID_JAG_TYP: Final = "oauth-id-jag+jwt" + + +@dataclass(frozen=True, slots=True) +class AssertionRejection: + """Why a JWT assertion failed verification. The reason is for the proxy log only, since + PyJWT's wording can quote the token's claims.""" + + reason: str + + class _JWTAuthSettings(Protocol): """The JWT auth settings block this handler reads back through ``getattr``, when one is configured.""" + @property + def id_jag_requesting_apps(self) -> Sequence[IdJagRequestingApp] | None: ... + @property def issuers(self) -> Sequence[JWTIssuerConfig] | None: ... @@ -282,7 +301,7 @@ class JWTHandler: return is_jwt_token(token=token) @staticmethod - def get_unverified_claims(token: str) -> dict | None: + def get_unverified_claims(token: str) -> dict[str, object] | None: """ Decode JWT claims without signature verification. Used for routing decisions before selecting validation path. @@ -949,7 +968,7 @@ class JWTHandler: keys: Final = response_json["keys"] if "keys" in response_json else response_json return cast(JWKKeyValue, keys) # cast-ok: JWTKeyItem declares only `kid`, validating would drop key material - async def _get_public_key_from_jwks_url(self, jwks_url: str, kid: str | None) -> dict: + async def _get_public_key_from_jwks_url(self, jwks_url: str, kid: str | None) -> JWTKeyItem: resolved_jwks_url: Final = await self._resolve_jwks_url(jwks_url) keys: Final = await self._cached_with_stale_fallback( cache_key=f"litellm_jwt_auth_keys_{resolved_jwks_url}", @@ -960,11 +979,11 @@ class JWTHandler: public_key: Final = self.parse_keys(keys=keys, kid=kid) if public_key is not None: - return cast(dict, public_key) + return public_key raise NoMatchingJWTPublicKeyError(f"No matching public key found. keys={resolved_jwks_url}, kid={kid}") - async def get_public_key(self, kid: str | None) -> dict: + async def get_public_key(self, kid: str | None) -> JWTKeyItem: keys_url: Final = os.getenv("JWT_PUBLIC_KEY_URL") if keys_url is None: @@ -1151,7 +1170,7 @@ class JWTHandler: # its jwks_uri, matching JWTIssuerConfig.jwks_url's documented fallback. return f"{issuer_config.issuer.rstrip('/')}/.well-known/openid-configuration" - def _get_claim_value_for_issuer_mapping(self, token: dict, claim_field: str) -> object: + def _get_claim_value_for_issuer_mapping(self, token: Mapping[str, object], claim_field: str) -> object: """Resolve a mapped claim from ``token``. Returns ``None`` when the field is absent or empty so that mapped claims @@ -1168,8 +1187,10 @@ class JWTHandler: return None return claim_value - def _apply_issuer_claim_mappings(self, token: dict, issuer_config: JWTIssuerConfig) -> dict: - normalized: Final[dict] = {k: v for k, v in token.items() if k not in self.LITELLM_INTERNAL_CLAIMS} + def _apply_issuer_claim_mappings( + self, token: Mapping[str, object], issuer_config: JWTIssuerConfig + ) -> dict[str, object]: + normalized: Final[dict[str, object]] = {k: v for k, v in token.items() if k not in self.LITELLM_INTERNAL_CLAIMS} normalized[self.LITELLM_JWT_ISSUER_CLAIM] = issuer_config.issuer claim_mappings: Final = [ (issuer_config.user_id_jwt_field, self.LITELLM_USER_ID_CLAIM), @@ -1192,8 +1213,8 @@ class JWTHandler: return normalized - def _get_jwk_from_public_key(self, public_key: dict) -> dict: - jwk: Final = {} + def _get_jwk_from_public_key(self, public_key: Mapping[str, object]) -> dict[str, object]: + jwk: Final[dict[str, object]] = {} for key in ["kty", "kid", "n", "e", "x", "y", "crv"]: if key in public_key: jwk[key] = public_key[key] @@ -1223,12 +1244,12 @@ class JWTHandler: def _decode_jwt_with_public_key( self, token: str, - public_key: dict | str, + public_key: JWTKeyItem | str, audience: str | list[str] | None, issuer: str | None = None, options: dict | None = None, disable_audience_validation: bool = False, - ) -> dict: + ) -> dict[str, object]: decode_options: Final = ( options if options is not None @@ -1298,11 +1319,64 @@ class JWTHandler: issuer_config=issuer_config, ) + async def decode_trusted_assertion( + self, token: str, audience: Sequence[str], required_claims: Sequence[str] + ) -> Mapping[str, object] | AssertionRejection: + """Verify a JWT assertion handed to the gateway's token endpoint with the issuer plumbing + JWT auth already trusts, a matching ``issuers`` entry or ``JWT_ISSUER`` with + ``JWT_PUBLIC_KEY_URL``, but against the caller's audience instead of the API audience. + An assertion whose ``iss`` nothing trusts is refused rather than checked unscoped. A JWKS + outage propagates as the 503 ``ProxyException`` so the caller does not call the token bad.""" + try: + header: Final = jwt.get_unverified_header(token) + except jwt.PyJWTError as malformed: + return AssertionRejection(reason=f"malformed assertion: {malformed}") + claims: Final = self.get_unverified_claims(token=token) + issuer_config: Final = self._get_configured_issuer(token=token) + trusted_issuer: Final = issuer_config.issuer if issuer_config is not None else os.getenv("JWT_ISSUER") + if not trusted_issuer or claims is None or claims.get("iss") != trusted_issuer: + return AssertionRejection(reason="iss is not an issuer this gateway trusts") + kid: Final = header.get("kid") + try: + public_key: Final = ( + await self._get_public_key_from_jwks_url(self._get_jwks_url_for_issuer(issuer_config), kid) + if issuer_config is not None + else await self.get_public_key(kid=kid) + ) + except JWKSUnreachableError as e: + raise jwks_unavailable_exception(e) from e + except Exception as no_key: # noqa: BLE001 # the key lookups raise plain Exceptions for a missing key + return AssertionRejection(reason=f"no signing key for the assertion: {no_key}") + try: + payload: Final = self._decode_jwt_with_public_key( + token=token, + public_key=public_key, + audience=list(audience), + issuer=trusted_issuer, + options={"require": list(required_claims)}, + ) + except jwt.PyJWTError as invalid: + return AssertionRejection(reason=f"assertion failed verification: {invalid}") + verified: Final = ( + payload + if issuer_config is None + else self._apply_issuer_claim_mappings(token=payload, issuer_config=issuer_config) + ) + validate: Final = self.litellm_jwtauth.custom_validate + if validate is not None and not validate(dict(verified)): + return AssertionRejection(reason="custom_validate rejected the assertion's claims") + return verified + async def auth_jwt(self, token: str) -> dict: header: Final = jwt.get_unverified_header(token) verbose_proxy_logger.debug("header: %s", header) + if header.get("typ") == ID_JAG_TYP: + raise IdJagIsNotABearerTokenError( + "an ID-JAG is an authorization grant, not a bearer token; redeem it at the token endpoint" + ) + kid: Final = header.get("kid", None) issuer_config: Final = self._get_configured_issuer(token=token) diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index ae129ca47b5..8dab693cd58 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -69,7 +69,8 @@ _PROXY_ADMIN_VIEW_ONLY_BLOCKED_ROUTES: Final = frozenset( # paths directly because the request route carries the resolved key id. _PROXY_ADMIN_VIEW_ONLY_BLOCKED_KEY_SUFFIXES: Final = ("/regenerate", "/reset_spend") -_AUTH_ENFORCED_PASS_THROUGH_ROUTE_GROUPS: Final = frozenset(("openai_routes", "llm_api_routes")) +_LLM_API_ROUTE_GROUPS: Final = frozenset(("llm_api_routes", "llm_api_routes_without_mcp")) +_AUTH_ENFORCED_PASS_THROUGH_ROUTE_GROUPS: Final = frozenset(("openai_routes", *_LLM_API_ROUTE_GROUPS)) # Method-scoped because the generic credential CRUD handlers share these paths: # the route check cannot admit PATCH/DELETE on a name like "user_connections" @@ -154,7 +155,7 @@ class RouteChecks: ################################################ # For llm_api_routes, also check registered pass-through endpoints ################################################ - if allowed_route == "llm_api_routes": + if allowed_route in _LLM_API_ROUTE_GROUPS: if route == "/auto_router/session" and RouteChecks.get_request_method(request) == "GET": return True @@ -185,7 +186,9 @@ class RouteChecks: # headers, env, credentials). POST/PUT/DELETE on # these paths are admin-only management writes and # are intentionally not covered. - if RouteChecks._is_get_mcp_server_discovery_route(route=route, request=request): + if allowed_route == "llm_api_routes" and RouteChecks._is_get_mcp_server_discovery_route( + route=route, request=request + ): return True # Agent registry CRUD moved from llm_api_routes into @@ -195,7 +198,7 @@ class RouteChecks: # could reach these paths, so keep them reachable here; # the handlers in agent_endpoints/endpoints.py still # enforce proxy-admin on writes and scope reads by role. - if RouteChecks.check_route_access( + if allowed_route == "llm_api_routes" and RouteChecks.check_route_access( route=route, allowed_routes=LiteLLMRoutes.agent_management_routes.value, ): diff --git a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index c96e759b1a3..c4414003ca5 100644 --- a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -4,6 +4,7 @@ import os from collections.abc import Sequence from collections.abc import Set as AbstractSet from datetime import datetime, timedelta, timezone +from types import SimpleNamespace from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch @@ -12,6 +13,7 @@ from fastapi import HTTPException from fastapi.testclient import TestClient from starlette.datastructures import Headers +from litellm.caching.caching import DualCache from litellm.constants import MCP_ALL_TOOLS_WILDCARD from litellm.models.access_group import LiteLLM_AccessGroupTable from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( @@ -23,6 +25,7 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( from litellm.proxy._types import ( LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable, + LiteLLM_UserTable, ProxyException, SpecialHeaders, SpecialMCPServerNames, @@ -11549,3 +11552,117 @@ class TestAccessGroupToolOwnershipAndFaults: patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader(groups, missing={"typo-group"})), ): assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == ["read"] + + +class _EmptyTables: + """Storage edge for a session credential's catalog refresh: the user row the real + ``reload_admitted_user`` reads, and every other table (keys, MCP servers, config) empty, so the + real ``_reload_admitted_key`` fails closed, since a session credential has no key row.""" + + def __init__(self, live_user: LiteLLM_UserTable) -> None: + self.litellm_usertable = SimpleNamespace(find_unique=AsyncMock(return_value=live_user)) + + def __getattr__(self, table: str) -> SimpleNamespace: + empty: Final = SimpleNamespace( + find_unique=AsyncMock(return_value=None), + find_first=AsyncMock(return_value=None), + find_many=AsyncMock(return_value=[]), + count=AsyncMock(return_value=0), + ) + setattr(self, table, empty) + return empty + + +def _session_credential_storage(live_user: LiteLLM_UserTable) -> SimpleNamespace: + tables: Final = _EmptyTables(live_user) + return SimpleNamespace(db=tables, writer_db=tables) + + +@pytest.mark.asyncio +async def test_catalog_refresh_reloads_session_credential_through_its_user_and_keeps_token_grants(monkeypatch): + """A ``/token``-minted or ``lite login`` credential has no key row: the catalog refresh must read the live + user policy, never the key table, and keep the team, routes and models sealed in the token.""" + from litellm.proxy import proxy_server + + database = _session_credential_storage( + LiteLLM_UserTable(user_id="session-user", organization_id="live-org", user_role="internal_user", teams=[]) + ) + monkeypatch.setattr(proxy_server, "prisma_client", database) + monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache()) + caller = UserAPIKeyAuth( + api_key="session-token-hash", + user_id="session-user", + team_id="token-team", + org_id="stale-org", + user_role="proxy_admin", + models=["token-model"], + allowed_routes=["mcp_inference_routes"], + is_session_token=True, + ) + caller.via_virtual_key = True + refreshed = await MCPRequestHandler.refresh_catalog_authority(caller) + assert database.writer_db.litellm_usertable.find_unique.await_count >= 1 + database.writer_db.litellm_verificationtoken.find_unique.assert_not_awaited() + assert (refreshed.object_permission, refreshed.object_permission_id) == (None, None) + assert (refreshed.user_role, refreshed.org_id) == ("internal_user", "live-org") + assert (refreshed.team_id, refreshed.models, refreshed.allowed_routes) == ( + "token-team", + ["token-model"], + ["mcp_inference_routes"], + ) + assert refreshed.is_session_token and refreshed.requires_fresh_policy + assert caller.org_id == "stale-org" and not caller.requires_fresh_policy + + +@pytest.mark.asyncio +@pytest.mark.parametrize("scope", ["litellm:mcp litellm:llm", "litellm:mcp", None]) +async def test_token_minted_credential_lists_tools_for_every_scope_shape(monkeypatch, scope): + """``tools/list`` with the credential ``/token`` mints (a session token the shared auth marks + ``via_virtual_key`` at the key-cache boundary) returns the tool list instead of the JSON-RPC + credential error, for the combined, MCP-only and unscoped ID-JAG shapes.""" + from mcp.types import ListToolsRequest, ListToolsResult, Tool + + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server import operations + from litellm.proxy._experimental.mcp_server.id_jag_grant import SubjectTokenRefusal, granted_scope + from litellm.proxy._experimental.mcp_server.operations import GatewayOperations, prepare_context + from litellm.proxy._types import MCPTransport + from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-session-credential-tools-list") + granted = granted_scope(scope, mcp_resource=False) + assert not isinstance(granted, SubjectTokenRefusal) + user = LiteLLM_UserTable(user_id="okta-user-1", user_email="okta@example.com", user_role="internal_user") + minted = ExperimentalUIJWTToken.get_cli_jwt_auth_token( + user, team_id="team-eng", team_alias="eng", allowed_routes=granted.allowed_routes + ) + credential = ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(minted) + assert credential is not None and credential.is_session_token + credential.api_key = "hash-of-" + minted[:16] + credential.via_virtual_key = True + server = MCPServer( + server_id="eng-server", + name="eng_server", + server_name="eng_server", + alias="eng_server", + url="https://upstream.example.com/mcp", + transport=MCPTransport.http, + ) + database = _session_credential_storage(user) + monkeypatch.setattr(proxy_server, "prisma_client", database) + monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache()) + tool = Tool(name="eng_server-list_repos", description="list repos", inputSchema={"type": "object"}) + with ( + patch.object(operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[server])), + patch.object( + operations.global_mcp_server_manager, + "get_tools_page", + AsyncMock(return_value=ListToolsResult(tools=[tool])), + ), + ): + result = await GatewayOperations().execute(ListToolsRequest(), prepare_context(user_api_key_auth=credential)) + assert [listed.name for listed in result.tools] == ["eng_server-list_repos"] + assert database.writer_db.litellm_usertable.find_unique.await_count >= 1 + database.writer_db.litellm_verificationtoken.find_unique.assert_not_awaited() + assert credential.team_id == "team-eng" and credential.allowed_routes == list(granted.allowed_routes or ()) diff --git a/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index f09409a023d..5fdc7669496 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -102,9 +102,8 @@ def isolate_global_mcp_registry(monkeypatch): restore, entries leak into other test modules sharing the same worker and break assertions over the full registry contents. """ - from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager - from litellm.proxy._experimental.mcp_server.catalog import CatalogSnapshots + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager monkeypatch.setattr(global_mcp_server_manager, "catalog", CatalogSnapshots(global_mcp_server_manager)) snapshot = dict(global_mcp_server_manager.registry) @@ -8133,9 +8132,9 @@ async def test_resolve_active_litellm_key_resolves_key_without_user_id(proxy_glo blocked and expiry, and the key hash (not the user) is what the mint seals. The per-user token store still gets no user for such a key, since there is none to key a stored credential by.""" from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( - extract_user_id_from_request, - _ResolvedKey, _resolve_active_litellm_key, + _ResolvedKey, + extract_user_id_from_request, ) from litellm.proxy._types import UserAPIKeyAuth, hash_token from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -8417,8 +8416,8 @@ async def test_load_active_user_by_id_serves_a_cached_row_without_a_database_rea cached row answers without a database read, and only a caller that asks for the database row pays for one.""" from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( - reload_active_user_by_id, load_active_user_by_id, + reload_active_user_by_id, ) from litellm.proxy._types import LiteLLM_UserTable from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -11428,11 +11427,10 @@ async def test_upstream_resource_sent_on_dcr_bridge_relay_authorize(): """The DCR-bridge relay arm builds its own upstream authorize params, so it needs the resource too. Without it the relayed authorize would omit the resource while the gateway's token leg still sent one, which is itself an invalid_target.""" - from litellm.types.mcp import MCPAuth - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( _dcr_bridge_relays_client_registration, ) + from litellm.types.mcp import MCPAuth server = _resource_server( client_id=None, @@ -11693,7 +11691,7 @@ def _native_client_app(monkeypatch): monkeypatch.setattr("litellm.proxy.proxy_server.master_key", _NATIVE_CLIENT_MASTER_KEY, raising=False) minted = [] - async def fake_mint(user_id, team_id): + async def fake_mint(user_id, team_id, allowed_routes=None): minted.append((user_id, team_id)) return MintedProxyCredential(key=f"sk-cli-{len(minted)}", expires_in=3600, user_id=user_id, team_id=team_id) @@ -12136,7 +12134,9 @@ def test_introspect_route_answers_for_authenticated_caller(monkeypatch): async def test_identity_bound_authorization_carries_nonce_and_caller_through_callback(monkeypatch): from http.cookies import SimpleCookie from urllib.parse import parse_qs, urlparse + from fastapi import Request + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( _oauth_state_cookie_name, authorize_with_server, @@ -12211,6 +12211,7 @@ async def test_identity_bound_authorization_carries_nonce_and_caller_through_cal @pytest.mark.asyncio async def test_enforced_login_warms_verified_token_readable_without_database_lookup(monkeypatch): from types import SimpleNamespace + from litellm.caching.caching import DualCache from litellm.proxy import proxy_server from litellm.proxy._experimental.mcp_server import db, mcp_server_manager @@ -12279,9 +12280,10 @@ async def test_enforced_login_warms_verified_token_readable_without_database_loo async def test_dcr_refusal_is_actionable_without_upstream_body( upstream_status: int, auth_type: MCPAuth, dcr_bridge: bool, flow: str, monkeypatch: pytest.MonkeyPatch ) -> None: - import httpx from typing import Final + import httpx + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( mint_ephemeral_dcr_client, register_client_with_server, @@ -12522,8 +12524,8 @@ async def test_oauth_exchange_stores_token_for_validated_jwt_user( import litellm from litellm.caching.llm_caching_handler import LLMClientCache from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler - from litellm.proxy import proxy_server from litellm.models.user import LiteLLM_UserTable + from litellm.proxy import proxy_server from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper from litellm.types.llms.custom_http import httpxSpecialProvider @@ -12618,8 +12620,8 @@ async def test_oauth_jwt_identity_rejects_untrusted_or_inactive_owner( from litellm.proxy import proxy_server from litellm.proxy._experimental.mcp_server import mcp_server_manager from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( - extract_user_id_from_request, authorize_oauth_credential_request, + extract_user_id_from_request, ) allowed_servers: Final = AsyncMock(return_value=["server-a"]) @@ -12688,7 +12690,7 @@ async def test_oauth_jwt_uses_configured_virtual_key_owner( ) -> None: from litellm.models.user import LiteLLM_UserTable from litellm.proxy._experimental.mcp_server.bridge_token_flow import extract_user_id_from_request - from litellm.proxy._types import UserAPIKeyAuth, UnregisteredJWTClientBehavior, hash_token + from litellm.proxy._types import UnregisteredJWTClientBehavior, UserAPIKeyAuth, hash_token from litellm.proxy.auth.auth_checks import jwt_key_mapping_cache_key handler, signing_key = jwt_oauth_identity @@ -13069,8 +13071,8 @@ async def test_signed_oauth_callback_honors_credential_write_policy( cimd: bool, ) -> None: import httpx - import litellm + import litellm from litellm.caching.llm_caching_handler import LLMClientCache from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy import proxy_server @@ -13186,17 +13188,16 @@ async def test_identity_bound_authorize_preserves_presented_jwt_permissions( allowed: bool, credential: str, ) -> None: - import jwt from datetime import datetime, timedelta, timezone from urllib.parse import parse_qs, urlparse + import jwt + from litellm.models.user import LiteLLM_UserTable from litellm.proxy import proxy_server - from litellm.proxy._types import JWTIssuerConfig, UserAPIKeyAuth, hash_token - from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken - from litellm.proxy._experimental.mcp_server import discoverable_endpoints, mcp_server_manager - from litellm.proxy._types import MCPTransport + from litellm.proxy._types import JWTIssuerConfig, MCPTransport, UserAPIKeyAuth, hash_token + from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer handler, signing_key = jwt_oauth_identity @@ -13330,9 +13331,10 @@ async def test_identity_bound_authorize_unrelated_bearer_uses_browser_session( credential: str, cookie_state: str, ) -> None: - import jwt from urllib.parse import parse_qs, urlparse + import jwt + from litellm.models.user import LiteLLM_UserTable from litellm.proxy import proxy_server from litellm.proxy._experimental.mcp_server import discoverable_endpoints, mcp_server_manager @@ -13664,6 +13666,7 @@ def test_persisted_client_is_not_applied_to_another_upstream(issuer, url): async def test_registration_does_not_write_client_after_upstream_edit(monkeypatch): monkeypatch.setenv("LITELLM_SALT_KEY", "test-mcp-registration-salt") from prisma import models + from litellm.proxy import proxy_server from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _persist_dcr_client_registration @@ -13750,6 +13753,7 @@ async def test_callback_does_not_forward_error_from_another_issuer(monkeypatch): @pytest.mark.parametrize("operation", ["authorize", "token"]) async def test_bound_client_cannot_be_sent_to_a_different_issuer(monkeypatch, operation): from fastapi import Request + from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -13796,6 +13800,7 @@ async def test_bound_client_cannot_be_sent_to_a_different_issuer(monkeypatch, op @pytest.mark.asyncio async def test_failed_registration_preserves_cached_credentials(monkeypatch): from fastapi import Request + from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints server = _dcr_redirect_test_server("old-client").model_copy( @@ -13819,8 +13824,13 @@ async def test_failed_registration_preserves_cached_credentials(monkeypatch): request.headers = {} with pytest.raises(HTTPException) as error: await endpoints.register_client_with_server( - request=request, mcp_server=server, client_name="app", grant_types=["authorization_code"], - response_types=["code"], token_endpoint_auth_method="none", replace_stale_client=True, + request=request, + mcp_server=server, + client_name="app", + grant_types=["authorization_code"], + response_types=["code"], + token_endpoint_auth_method="none", + replace_stale_client=True, ) assert error.value.status_code == 503 assert "could not be saved" in error.value.detail @@ -13835,6 +13845,7 @@ async def test_failed_registration_preserves_cached_credentials(monkeypatch): @pytest.mark.asyncio async def test_registration_without_database_keeps_client_in_temporary_server(monkeypatch): from fastapi import Request + from litellm.proxy import proxy_server from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints @@ -13850,8 +13861,13 @@ async def test_registration_without_database_keeps_client_in_temporary_server(mo request.base_url = "https://gateway.example/" request.headers = {} response = await endpoints.register_client_with_server( - request=request, mcp_server=server, client_name="app", grant_types=["authorization_code"], - response_types=["code"], token_endpoint_auth_method="none", replace_stale_client=True, + request=request, + mcp_server=server, + client_name="app", + grant_types=["authorization_code"], + response_types=["code"], + token_endpoint_auth_method="none", + replace_stale_client=True, ) assert response.status_code == 200 assert json.loads(response.body)["client_id"] == "temporary-client" @@ -13864,9 +13880,10 @@ async def test_registration_without_database_keeps_client_in_temporary_server(mo @pytest.mark.asyncio @pytest.mark.parametrize("config_store", [False, True]) async def test_registration_does_not_overwrite_credentials_after_failed_identity_read(monkeypatch, config_store): - from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints, db - from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager from litellm.proxy import utils + from litellm.proxy._experimental.mcp_server import db + from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager server = _dcr_redirect_test_server(None) monkeypatch.setattr(utils, "get_prisma_client_or_throw", lambda _: MagicMock()) @@ -13895,9 +13912,10 @@ async def test_registration_does_not_overwrite_credentials_after_failed_identity @pytest.mark.asyncio @pytest.mark.parametrize("winner_available", [False, True]) async def test_registration_losing_conditional_write_reuses_only_a_matching_winner(monkeypatch, winner_available): - from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints, db - from litellm.proxy._types import LiteLLM_MCPServerTable from litellm.proxy import utils + from litellm.proxy._experimental.mcp_server import db + from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints + from litellm.proxy._types import LiteLLM_MCPServerTable server = _dcr_redirect_test_server(None).model_copy(update={"url": "https://upstream.example/mcp"}) row = LiteLLM_MCPServerTable.model_validate(server.model_dump(exclude_none=True)) @@ -14503,7 +14521,8 @@ async def test_managed_authorize_retains_the_callers_previous_grant( from urllib.parse import parse_qs, urlsplit from litellm.proxy import proxy_server - from litellm.proxy._experimental.mcp_server import db, discoverable_endpoints as endpoints + from litellm.proxy._experimental.mcp_server import db + from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager _, signing_key = jwt_oauth_identity @@ -14548,7 +14567,8 @@ async def test_managed_code_exchange_preserves_requested_scope_only_when_omitted from urllib.parse import parse_qs, urlsplit from litellm.proxy import proxy_server - from litellm.proxy._experimental.mcp_server import db, discoverable_endpoints as endpoints + from litellm.proxy._experimental.mcp_server import db + from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager _, signing_key = jwt_oauth_identity @@ -14692,7 +14712,8 @@ async def test_managed_refresh_retains_only_the_actual_grant( from types import SimpleNamespace from litellm.proxy import proxy_server - from litellm.proxy._experimental.mcp_server import db, discoverable_endpoints as endpoints + from litellm.proxy._experimental.mcp_server import db + from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager _, signing_key = jwt_oauth_identity @@ -14774,7 +14795,8 @@ async def test_managed_resource_metadata_excludes_refresh_permission(monkeypatch ], ) async def test_connect_requires_the_requested_vendor_grant(monkeypatch, granted, required, expected): - from litellm.proxy._experimental.mcp_server import db, discoverable_endpoints as endpoints + from litellm.proxy._experimental.mcp_server import db + from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints credential = {**_stored_grant(), "scopes": granted} read = AsyncMock(return_value=credential) @@ -14836,3 +14858,137 @@ def test_aggregate_protected_resource_advertises_header_server_scopes( ) with patch.object(mcp_server_manager, "global_mcp_server_manager", manager): assert _build_aggregate_protected_resource_response(request)["scopes_supported"] == expected + + +@pytest.mark.asyncio +async def test_token_endpoint_routes_the_jwt_bearer_grant_to_the_aggregate_gateway(): + """A requesting app presents its identity provider's client_id, never a gateway DCR one, so + the aggregate dispatch keys on the grant type for RFC 7523 and hands the assertion to the + ID-JAG redeemer. A per-server token endpoint never serves that grant.""" + from litellm.proxy._experimental.mcp_server import discoverable_endpoints + from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import JWT_BEARER_GRANT_TYPE + from litellm.proxy._experimental.mcp_server.id_jag_grant import IdJagPrerequisites + + request = _mock_callback_request("https://litellm.example.com/") + with patch("litellm.proxy.proxy_server.master_key", "sk-test-salt-for-lit9670"): + response = await discoverable_endpoints.token_endpoint( + request=request, + grant_type=JWT_BEARER_GRANT_TYPE, + client_id="0oa-requesting-app", + client_secret="app-secret", + assertion="eyJ.assertion.sig", + resource=None, + ) + assert response.status_code == 400 + refused = json.loads(bytes(response.body)) + assert refused["error"] == "unsupported_grant_type" + unmet = IdJagPrerequisites( + jwt_auth_enabled=False, + has_database=False, + licensed=False, + maps_jwts_to_virtual_keys=False, + shared_replay_store=False, + base_url=None, + requesting_apps=(), + ).refusal() + assert unmet is not None and refused["error_description"] == unmet.description + with pytest.raises(HTTPException): + await discoverable_endpoints.token_endpoint( + request=request, + grant_type=JWT_BEARER_GRANT_TYPE, + client_id="0oa-requesting-app", + assertion="eyJ.assertion.sig", + mcp_server_name="no-such-server", + ) + + +@pytest.mark.asyncio +async def test_token_endpoint_takes_the_client_id_from_http_basic_when_the_form_leaves_it_out(): + """RFC 6749 section 2.3.1: a confidential app may authenticate with HTTP Basic alone, so the + aggregate dispatch reads client_id from that header. With neither the form field nor the + header the request is a 400 invalid_request, not a 422.""" + from base64 import b64encode + + from litellm.proxy._experimental.mcp_server import discoverable_endpoints + from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import JWT_BEARER_GRANT_TYPE + + request = _mock_callback_request("https://litellm.example.com/") + request.headers = {"authorization": "Basic " + b64encode(b"0oa-requesting-app:app-secret").decode()} + with patch("litellm.proxy.proxy_server.master_key", "sk-test-salt-for-lit9670"): + response = await discoverable_endpoints.token_endpoint( + request=request, + grant_type=JWT_BEARER_GRANT_TYPE, + client_id=None, + assertion="eyJ.assertion.sig", + resource=None, + ) + assert response.status_code == 400 + assert json.loads(bytes(response.body))["error"] == "unsupported_grant_type" + request.headers = {} + with pytest.raises(HTTPException) as refused: + await discoverable_endpoints.token_endpoint( + request=request, grant_type=JWT_BEARER_GRANT_TYPE, client_id=None, assertion="eyJ.assertion.sig" + ) + assert refused.value.status_code == 400 + assert refused.value.detail["error"] == "invalid_request" + + +@pytest.mark.parametrize("available", [True, False]) +def test_root_metadata_is_the_gateways_own_and_lists_the_jwt_bearer_grant_when_it_is_served(available): + """With no OAuth2 MCP server resolving at the root, the bare-origin document is the gateway's + own authorization server metadata: issuer is the origin an ID-JAG names as ``aud``, the token + endpoint is ``{base}/token``, and the jwt-bearer grant (with HTTP Basic client auth) is listed + exactly when the gateway serves it, under the same gates as the aggregate ``/mcp`` document.""" + from fastapi import Request + + from litellm.proxy._experimental.mcp_server import discoverable_endpoints + from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import JWT_BEARER_GRANT_TYPE + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + global_mcp_server_manager.registry.clear() + request = MagicMock(spec=Request) + request.base_url = "https://llm.example.com/" + request.headers = {} + with ( + patch.object(discoverable_endpoints, "jwt_bearer_available", return_value=available), # test-quality-ok: deployment gate seam + patch.object(discoverable_endpoints, "token_exchange_available", return_value=True), # test-quality-ok: deployment gate seam + ): + document = discoverable_endpoints._build_oauth_authorization_server_response(request=request, mcp_server_name=None) + assert document["issuer"] == "https://llm.example.com" + assert document["token_endpoint"] == "https://llm.example.com/token" + assert (JWT_BEARER_GRANT_TYPE in document["grant_types_supported"]) is available + assert ("client_secret_basic" in document["token_endpoint_auth_methods_supported"]) is available + assert ("none" in document["token_endpoint_auth_methods_supported"]) is available + assert "urn:ietf:params:oauth:grant-type:token-exchange" in document["grant_types_supported"] + + +def test_root_metadata_owned_by_a_resolved_oauth2_server_never_lists_the_gateways_grant(): + """When exactly one OAuth2 MCP server resolves at the root, the root document describes that + upstream's token endpoint, which cannot redeem an ID-JAG, so the jwt-bearer grant is not + advertised there even though the gateway serves it; a requesting app uses the aggregate + ``/.well-known/oauth-authorization-server/mcp`` document in that deployment.""" + from fastapi import Request + + from litellm.proxy._experimental.mcp_server import discoverable_endpoints + from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import JWT_BEARER_GRANT_TYPE + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + server = _create_oauth2_server() + global_mcp_server_manager.registry.clear() + global_mcp_server_manager.registry[server.server_id] = server + request = MagicMock(spec=Request) + request.base_url = "https://llm.example.com/" + request.headers = {} + try: + with ( + patch.object(discoverable_endpoints, "jwt_bearer_available", return_value=True), # test-quality-ok: deployment gate seam + patch.object(discoverable_endpoints, "token_exchange_available", return_value=True), # test-quality-ok: deployment gate seam + patch.object(global_mcp_server_manager, "get_mcp_server_by_name", return_value=server), # test-quality-ok: resolver seam + ): + document = discoverable_endpoints._build_oauth_authorization_server_response(request=request, mcp_server_name=None) + finally: + global_mcp_server_manager.registry.clear() + assert document["token_endpoint"] == "https://llm.example.com/test_oauth/token" + assert JWT_BEARER_GRANT_TYPE not in document["grant_types_supported"] + assert document["grant_types_supported"] == ["authorization_code", "refresh_token"] + assert document["token_endpoint_auth_methods_supported"] == ["client_secret_post"] diff --git a/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py b/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py index 1417e77c8e2..ef95df23ff8 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py @@ -20,10 +20,12 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( CONNECT_FLOW_COOKIE_PREFIX, GATEWAY_AUTH_CODE_PREFIX, GATEWAY_AUTH_CODE_TTL_SECONDS, + JWT_BEARER_GRANT_TYPE, MANUAL_DELIVERY_AUTH_CODE_TTL_SECONDS, MAX_CLIENT_ID_LENGTH, SUBJECT_TOKEN_TYPES, TOKEN_EXCHANGE_GRANT_TYPE, + AssertionIdentity, ConsentTeam, MintedProxyCredential, SubjectIdentity, @@ -44,6 +46,7 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( open_gateway_dcr_client, register_aggregate_client, revoke_refresh_token, + supported_grant_types, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import ( SessionBearerAdmitted, @@ -1771,10 +1774,12 @@ CONSENT_TEAMS = (ConsentTeam(team_id="team-a", team_alias="Team A"), ConsentTeam class _Minter: def __init__(self, result=None): self.calls = [] + self.allowed_routes = [] self.result = result - async def __call__(self, user_id, team_id): + async def __call__(self, user_id, team_id, allowed_routes=None): self.calls.append((user_id, team_id)) + self.allowed_routes.append(allowed_routes) if self.result is not None: return self.result return MintedProxyCredential(key=f"sk-cli-{user_id}", expires_in=3600, user_id=user_id, team_id=team_id) @@ -3296,3 +3301,200 @@ async def test_gateway_scope_still_narrows_when_the_resource_does_not_resolve() _, cookies = _flow_cookie_from(response) sealed_flow: Final = _sealed_wire_json(next(iter(cookies.values())), "", "gateway_connect_flow") assert sealed_flow["resource_server_id"] == "alpha-id" + + +ID_JAG = "eyJhbGciOiJSUzI1NiIsInR5cCI6Im9hdXRoLWlkLWphZytqd3QifQ.eyJzdWIiOiJ1MSJ9.idp-signature" +REQUESTING_APP = "0oa-requesting-app" +ASSERTED = AssertionIdentity( + user_id="u1", + team_id="team-b", + replay_key="mcp_gateway_dcr:id_jag:jti-1", + replay_ttl_seconds=300, + scope="litellm:mcp", + allowed_routes=("mcp_inference_routes",), +) + + +class _Redeemer: + def __init__(self, result=None): + self.calls = [] + self.result = result + + async def __call__(self, assertion, client_id, client_secret, request): + self.calls.append((assertion, client_id, client_secret, request.url.path)) + return self.result if self.result is not None else ASSERTED + + +async def _redeem_assertion_native(minter, redeemer, cache=None, client_id=REQUESTING_APP, **overrides): + arguments = { + "grant_type": JWT_BEARER_GRANT_TYPE, + "assertion": ID_JAG, + "client_secret": "app-secret", + "redeem_assertion": redeemer, + } + return await _redeem_native(None, client_id, minter, cache=cache, **{**arguments, **overrides}) + + +@pytest.mark.asyncio +async def test_jwt_bearer_grant_mints_the_proxy_credential_for_the_asserted_user_once_per_assertion(): + """RFC 7523: an allowlisted requesting app redeems the ID-JAG its identity provider minted for + the same credential the consent flow mints, attributed to the user the assertion names and + the team asked for, with the assertion's scope stated and no refresh token (RFC 7521 section + 4.1, the app renews with a fresh assertion). The assertion's jti is single-use: presenting it again is a replay and mints + nothing, while a fresh assertion for the same user is served.""" + minter, redeemer, cache = _Minter(), _Redeemer(), DualCache() + response = await _redeem_assertion_native(minter, redeemer, cache=cache) + assert response.status_code == 200 + assert response.headers["cache-control"] == "no-store" + body = json.loads(response.body) + assert redeemer.calls == [(ID_JAG, REQUESTING_APP, "app-secret", "/token")] + assert minter.calls == [("u1", "team-b")] + assert body["access_token"] == "sk-cli-u1" + assert body["token_type"] == "Bearer" + assert body["scope"] == "litellm:mcp" + assert "issued_token_type" not in body + assert (body["user_id"], body["team_id"]) == ("u1", "team-b") + assert "refresh_token" not in body + replayed = await _redeem_assertion_native(minter, redeemer, cache=cache) + assert replayed.status_code == 400 + assert json.loads(replayed.body)["error"] == "invalid_grant" + assert minter.calls == [("u1", "team-b")] + fresh = _Redeemer(ASSERTED.model_copy(update={"replay_key": "mcp_gateway_dcr:id_jag:jti-2", "scope": None})) + again = await _redeem_assertion_native(minter, fresh, cache=cache) + assert again.status_code == 200 + assert "scope" not in json.loads(again.body) + assert minter.calls == [("u1", "team-b"), ("u1", "team-b")] + + +@pytest.mark.asyncio +async def test_jwt_bearer_grant_claims_the_jti_before_it_mints(): + """A mint the gateway could not complete still consumed the assertion: the client gets the + 503 and comes back with a fresh assertion, never with the same one.""" + minter, cache = _Minter(result="unavailable"), DualCache() + first = await _redeem_assertion_native(minter, _Redeemer(), cache=cache) + assert first.status_code == 503 + retried = await _redeem_assertion_native(_Minter(), _Redeemer(), cache=cache) + assert retried.status_code == 400 + assert json.loads(retried.body)["error"] == "invalid_grant" + + +@pytest.mark.parametrize( + "error, status", + [ + ("invalid_client", 401), + ("invalid_grant", 400), + ("invalid_request", 400), + ("unsupported_grant_type", 400), + ("temporarily_unavailable", 503), + ], +) +@pytest.mark.asyncio +async def test_jwt_bearer_grant_refusals_keep_their_oauth_error_and_status(error, status): + minter, cache = _Minter(), DualCache() + refusal = SubjectTokenRefusal(error=error, description="why") + response = await _redeem_assertion_native(minter, _Redeemer(refusal), cache=cache) + assert response.status_code == status + assert json.loads(response.body) == {"error": error, "error_description": "why"} + assert minter.calls == [] + served = await _redeem_assertion_native(minter, _Redeemer(), cache=cache) + assert served.status_code == 200 + + +@pytest.mark.asyncio +async def test_jwt_bearer_grant_requires_an_assertion_and_the_proxy_api_target(): + minter, redeemer = _Minter(), _Redeemer() + missing = await _redeem_assertion_native(minter, redeemer, assertion=None) + assert missing.status_code == 400 + assert json.loads(missing.body)["error"] == "invalid_request" + foreign = await _redeem_assertion_native(minter, redeemer, resource="https://elsewhere.example/api") + assert foreign.status_code == 400 + assert json.loads(foreign.body)["error"] == "invalid_target" + assert redeemer.calls == [] + assert minter.calls == [] + + +@pytest.mark.asyncio +async def test_jwt_bearer_grant_is_refused_where_no_redeemer_is_wired(): + minter = _Minter() + response = await _redeem_native( + None, REQUESTING_APP, minter, grant_type=JWT_BEARER_GRANT_TYPE, assertion=ID_JAG, client_secret=None + ) + assert response.status_code == 400 + assert json.loads(response.body)["error"] == "unsupported_grant_type" + assert minter.calls == [] + + +@pytest.mark.asyncio +async def test_an_unknown_grant_type_names_the_jwt_bearer_grant_among_the_served_ones(): + response = await _redeem_native(None, REQUESTING_APP, _Minter(), grant_type="urn:example:nope") + assert response.status_code == 400 + assert JWT_BEARER_GRANT_TYPE in json.loads(response.body)["error_description"] + + +@pytest.mark.parametrize( + "token_exchange, jwt_bearer, expected", + [ + (False, False, ("authorization_code", "refresh_token")), + (True, False, ("authorization_code", "refresh_token", TOKEN_EXCHANGE_GRANT_TYPE)), + (False, True, ("authorization_code", "refresh_token", JWT_BEARER_GRANT_TYPE)), + (True, True, ("authorization_code", "refresh_token", TOKEN_EXCHANGE_GRANT_TYPE, JWT_BEARER_GRANT_TYPE)), + ], +) +def test_the_jwt_bearer_grant_is_advertised_exactly_when_it_is_served(token_exchange, jwt_bearer, expected): + assert supported_grant_types(token_exchange, jwt_bearer) == expected + contract = native_client_auth_contract(_request("/.well-known/litellm-cli-auth"), token_exchange, jwt_bearer) + assert contract["grant_types_supported"] == expected + + +@pytest.mark.asyncio +async def test_jwt_bearer_grant_narrows_the_credential_to_the_routes_the_scope_bought(): + """The route groups the assertion's scope resolved to reach the minter, so the credential is + enforced on every proxy route the way a virtual key's allowed_routes is, and the response states + exactly the scope that was enforced. No scope on the assertion mints the full credential.""" + minter = _Minter() + both = ASSERTED.model_copy( + update={"scope": "litellm:mcp litellm:llm", "allowed_routes": ("mcp_inference_routes", "openai_routes")} + ) + scoped = await _redeem_assertion_native(minter, _Redeemer(both), cache=DualCache()) + assert scoped.status_code == 200 + assert json.loads(scoped.body)["scope"] == "litellm:mcp litellm:llm" + full = await _redeem_assertion_native( + minter, _Redeemer(ASSERTED.model_copy(update={"scope": None, "allowed_routes": None})), cache=DualCache() + ) + assert full.status_code == 200 + assert "scope" not in json.loads(full.body) + assert minter.allowed_routes == [("mcp_inference_routes", "openai_routes"), None] + + +@pytest.mark.asyncio +async def test_jwt_bearer_grant_accepts_the_mcp_resource_indicator_as_its_target(): + """RFC 8707 clients send the ID-JAG's audience as ``resource``, and that audience may be the MCP + resource ``{base}/mcp``, so it is a valid target next to the proxy origin. Anything else stays + ``invalid_target``.""" + minter = _Minter() + served = await _redeem_assertion_native( + minter, _Redeemer(), cache=DualCache(), resource=PROXY_API_RESOURCE.rstrip("/") + "/mcp" + ) + assert served.status_code == 200 + foreign = await _redeem_assertion_native( + minter, _Redeemer(), cache=DualCache(), resource=PROXY_API_RESOURCE.rstrip("/") + "/other/mcp" + ) + assert json.loads(foreign.body)["error"] == "invalid_target" + assert minter.calls == [("u1", "team-b")] + + +@pytest.mark.parametrize("jwt_bearer_available", [True, False]) +def test_aggregate_metadata_advertises_basic_client_auth_only_with_the_grant_that_reads_it(jwt_bearer_available): + """Only the jwt-bearer redemption reads an HTTP Basic client secret, the other grants take + ``client_secret`` from the form, so ``client_secret_basic`` is advertised only when that grant + is on. A client following the document otherwise sends Basic alone and authenticates with no + secret.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _build_aggregate_authorization_server_response, + ) + + document = _build_aggregate_authorization_server_response(_request(), True, jwt_bearer_available) + methods = document["token_endpoint_auth_methods_supported"] + assert isinstance(methods, list) + assert ("client_secret_basic" in methods) is jwt_bearer_available + assert "none" in methods and "client_secret_post" in methods diff --git a/tests/unit/proxy/_experimental/mcp_server/test_id_jag_grant.py b/tests/unit/proxy/_experimental/mcp_server/test_id_jag_grant.py new file mode 100644 index 00000000000..fdfa0544ad1 --- /dev/null +++ b/tests/unit/proxy/_experimental/mcp_server/test_id_jag_grant.py @@ -0,0 +1,738 @@ +import hashlib +import logging +import math +from base64 import b64encode +from collections.abc import Mapping, Sequence +from typing import Final + +import jwt +import pytest +from fastapi import HTTPException + +from litellm.caching.caching import DualCache +from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( + AssertionIdentity, + ConsentTeam, + ReloadUserFailure, + SubjectIdentity, + SubjectTokenRefusal, +) +from litellm.proxy._experimental.mcp_server.id_jag_grant import ( + ID_JAG_REQUIRED_CLAIMS, + REJECTED_ASSERTION, + TEAM_HEADER, + AssertionRedemption, + GrantedScope, + IdJagPrerequisites, + LookupTeams, + _claim_reader, + basic_client_id, + gateway_audiences, + granted_scope, + identity_from_id_jag, + jwt_bearer_available, + names_mcp_resource, + proxy_user_lookup, + select_team, + subject_resolver, +) +from litellm.proxy._types import IdJagRequestingApp, LiteLLM_JWTAuth, LiteLLMRoutes +from litellm.proxy.auth.handle_jwt import ( + ID_JAG_TYP, + AssertionRejection, + JWKSUnreachableError, + JWTHandler, + jwks_unavailable_exception, +) + +ISSUER = "https://idp.example.com" +GATEWAY = "https://gateway.example.com" +AUDIENCES = (f"{GATEWAY}/mcp", GATEWAY) +NOW = 1_800_000_000.0 +PUBLIC_APP = IdJagRequestingApp(client_id="0oa-public") +CONFIDENTIAL_APP = IdJagRequestingApp(client_id="0oa-confidential", client_secret="s3cret") +SPACED_APP = IdJagRequestingApp(client_id="0oa spaced", client_secret="app secret") +EVERY_GATE_HOLDS = { + "jwt_auth_enabled": True, + "has_database": True, + "licensed": True, + "maps_jwts_to_virtual_keys": False, + "shared_replay_store": True, + "base_url": "https://gateway.example.com", + "requesting_apps": (PUBLIC_APP, CONFIDENTIAL_APP, SPACED_APP), +} +JWKS_DOWN = jwks_unavailable_exception(JWKSUnreachableError("ConnectError fetching the IdP's JWKS after 3 attempts")) + + +def _claims(**overrides: object) -> dict[str, object]: + return { + "iss": ISSUER, + "sub": "okta-user-1", + "aud": AUDIENCES[1], + "exp": NOW + 300, + "iat": NOW - 5, + "jti": "jti-1", + "client_id": PUBLIC_APP.client_id, + "scope": "litellm:mcp", + **overrides, + } + + +def _assertion(typ: str = ID_JAG_TYP, **claims: object) -> str: + return jwt.encode( + _claims(**claims), + "unused-by-the-fake-verifier-but-long-enough-for-hs256", + algorithm="HS256", + headers={"typ": typ}, + ) + + +class _Verifier: + """Stands in for the injected ``DecodeAssertion`` port; records each call as a new tuple.""" + + def __init__( + self, result: Mapping[str, object] | AssertionRejection | None = None, raises: BaseException | None = None + ) -> None: + self.calls: tuple[tuple[str, tuple[str, ...], tuple[str, ...]], ...] = () + self.result = result + self.raises = raises + + async def __call__( + self, assertion: str, audience: Sequence[str], required_claims: Sequence[str] + ) -> Mapping[str, object] | AssertionRejection: + self.calls = (*self.calls, (assertion, tuple(audience), tuple(required_claims))) # rebind-ok: immutable record + if self.raises is not None: + raise self.raises + return self.result if self.result is not None else jwt.decode(assertion, options={"verify_signature": False}) + + +class _Resolver: + """Stands in for the injected ``ResolveSubject`` port; records each call as a new tuple.""" + + def __init__( + self, + user_id: str | None = "u1", + raises: BaseException | None = None, + result: SubjectIdentity | SubjectTokenRefusal | None = None, + ) -> None: + self.calls: tuple[tuple[Mapping[str, object], str | None], ...] = () + self.user_id = user_id + self.raises = raises + self.result = result + + async def __call__( + self, claims: Mapping[str, object], team_header: str | None + ) -> SubjectIdentity | SubjectTokenRefusal | None: + self.calls = (*self.calls, (dict(claims), team_header)) # rebind-ok: immutable record + if self.raises is not None: + raise self.raises + if self.result is not None: + return self.result + return None if self.user_id is None else SubjectIdentity(user_id=self.user_id, team_id=team_header) + + +def _redemption( + assertion: str, + client_id: str = PUBLIC_APP.client_id, + client_secret: str | None = None, + authorization: str | None = None, + team_id: str | None = "team-b", + leeway: float = 0.0, +) -> AssertionRedemption: + return AssertionRedemption( + assertion=assertion, + client_id=client_id, + client_secret=client_secret, + authorization_header=authorization, + team_id=team_id, + audiences=AUDIENCES, + now=NOW, + leeway_seconds=leeway, + ) + + +@pytest.mark.asyncio +async def test_the_replay_marker_outlives_the_leeway_the_verifier_grants_exp(): + """PyJWT accepts an assertion until ``exp`` plus the handler's leeway, so the marker must + live that long too, including for an assertion already past ``exp`` but inside the leeway.""" + within_leeway = await _identity(_redemption(_assertion(), leeway=30), _Verifier(), _Resolver()) + assert isinstance(within_leeway, AssertionIdentity) + assert within_leeway.replay_ttl_seconds == 330 + just_expired = await _identity( + _redemption(_assertion(), leeway=30), _Verifier(result=_claims(exp=NOW - 10)), _Resolver() + ) + assert isinstance(just_expired, AssertionIdentity) + assert just_expired.replay_ttl_seconds == 20 + + +async def _identity( + redemption: AssertionRedemption, + verifier: _Verifier | None = None, + resolver: _Resolver | None = None, + **unmet: object, +) -> AssertionIdentity | SubjectTokenRefusal: + return await identity_from_id_jag( + redemption, + prerequisites=IdJagPrerequisites(**{**EVERY_GATE_HOLDS, **unmet}), + decode=verifier if verifier is not None else _Verifier(), + resolve_subject=resolver if resolver is not None else _Resolver(), + ) + + +@pytest.mark.asyncio +async def test_a_valid_id_jag_names_the_user_its_single_use_key_and_its_scope(): + """The verifier is asked for the gateway's own issuer values as audience and for every claim + the ID-JAG profile requires, the subject lookup sees the verified claims, and the identity + carries the team asked for, a replay key derived from iss and jti, a marker TTL that outlives + the assertion, and the assertion's scope.""" + assertion, verifier, resolver = _assertion(), _Verifier(), _Resolver() + identity = await _identity(_redemption(assertion), verifier, resolver) + assert isinstance(identity, AssertionIdentity) + assert verifier.calls == ((assertion, AUDIENCES, ID_JAG_REQUIRED_CLAIMS),) + assert {"iss", "sub", "aud", "exp", "iat", "jti", "client_id"} <= set(ID_JAG_REQUIRED_CLAIMS) + assert resolver.calls == ((_claims(), "team-b"),) + assert (identity.user_id, identity.team_id, identity.scope) == ("u1", "team-b", "litellm:mcp") + assert identity.allowed_routes == ("mcp_inference_routes",) + assert identity.replay_key == "mcp_gateway_dcr:id_jag:" + hashlib.sha256(f"{ISSUER}\njti-1".encode()).hexdigest() + assert identity.replay_ttl_seconds == math.ceil(300) + other_jti = await _identity(_redemption(_assertion(jti="jti-2")), _Verifier(), _Resolver()) + assert isinstance(other_jti, AssertionIdentity) + assert other_jti.replay_key != identity.replay_key + teamless = await _identity(_redemption(_assertion(scope=""), team_id=None)) + assert isinstance(teamless, AssertionIdentity) + assert (teamless.team_id, teamless.scope, teamless.allowed_routes) == (None, None, None) + + +@pytest.mark.asyncio +async def test_a_requesting_app_the_operator_did_not_allowlist_learns_nothing_about_the_assertion(): + verifier, resolver = _Verifier(), _Resolver() + refused = await _identity(_redemption(_assertion(), client_id="0oa-stranger"), verifier, resolver) + assert refused == SubjectTokenRefusal( + error="invalid_client", description="client_id is not an allowlisted requesting app" + ) + assert verifier.calls == () + assert resolver.calls == () + + +def _basic(client_id: str, secret: str) -> str: + return "Basic " + b64encode(f"{client_id}:{secret}".encode()).decode() + + +@pytest.mark.parametrize( + "client_secret, authorization, accepted", + [ + (None, None, False), + ("wrong", None, False), + ("s3cr\u00e9t", None, False), + ("s3cret", None, True), + (None, _basic(CONFIDENTIAL_APP.client_id, "s3cr\u00e9t"), False), + (None, _basic(CONFIDENTIAL_APP.client_id, "s3cret"), True), + (None, _basic(CONFIDENTIAL_APP.client_id, "wrong"), False), + (None, _basic(PUBLIC_APP.client_id, "s3cret"), False), + ("s3cret", _basic(CONFIDENTIAL_APP.client_id, "wrong"), False), + (None, "Basic not-base64!", False), + (None, "Bearer s3cret", False), + ], +) +@pytest.mark.asyncio +async def test_a_confidential_app_must_authenticate_with_its_secret(client_secret, authorization, accepted): + verifier = _Verifier() + assertion = _assertion(client_id=CONFIDENTIAL_APP.client_id) + redemption = _redemption(assertion, CONFIDENTIAL_APP.client_id, client_secret, authorization) + identity = await _identity(redemption, verifier) + if accepted: + assert isinstance(identity, AssertionIdentity) + return + assert identity == SubjectTokenRefusal(error="invalid_client", description="client authentication failed") + assert verifier.calls == () + + +@pytest.mark.parametrize( + "client_secret, authorization, accepted", + [ + ("app secret", None, True), + (None, "Basic " + b64encode(b"0oa+spaced:app+secret").decode(), True), + (None, "Basic " + b64encode(b"0oa%20spaced:app%20secret").decode(), True), + (None, "Basic " + b64encode(b"0oa spaced:app secret").decode(), True), + (None, "Basic " + b64encode(b"0oa+spaced:app+wrong").decode(), False), + (None, "Basic " + b64encode(b"0oa-spaced:app+secret").decode(), False), + ], +) +@pytest.mark.asyncio +async def test_basic_credentials_are_form_decoded_so_a_space_in_an_id_or_secret_still_matches( + client_secret, authorization, accepted +): + """RFC 6749 section 2.3.1 form-encodes the id and secret before base64, so ``+`` is a space. The + endpoint hands over the ``client_id`` it read from the header, so the header must name the same app.""" + assertion = _assertion(client_id=SPACED_APP.client_id) + identity = await _identity(_redemption(assertion, SPACED_APP.client_id, client_secret, authorization), _Verifier()) + if accepted: + assert isinstance(identity, AssertionIdentity) + return + assert identity == SubjectTokenRefusal(error="invalid_client", description="client authentication failed") + + +@pytest.mark.asyncio +async def test_a_public_app_is_identified_by_client_id_alone(): + identity = await _identity(_redemption(_assertion(), client_secret="anything", authorization="Bearer x")) + assert isinstance(identity, AssertionIdentity) + + +@pytest.mark.parametrize( + "assertion, description", + [ + ("not-a-jwt", "assertion is not a JWT"), + (_assertion(typ="JWT"), f"assertion typ must be {ID_JAG_TYP}"), + (_assertion(typ="at+jwt"), f"assertion typ must be {ID_JAG_TYP}"), + ], +) +@pytest.mark.asyncio +async def test_a_bearer_or_id_token_is_not_an_id_jag(assertion, description): + verifier = _Verifier() + refused = await _identity(_redemption(assertion), verifier) + assert refused == SubjectTokenRefusal(error="invalid_grant", description=description) + assert verifier.calls == () + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "verifier, reason", + [ + (_Verifier(result=AssertionRejection(reason="Signature verification failed")), "Signature verification"), + (_Verifier(result=AssertionRejection(reason="Audience doesn't match")), "Audience"), + (_Verifier(result=AssertionRejection(reason='Token is missing the "jti" claim')), "jti"), + (_Verifier(result=AssertionRejection(reason="Signature has expired")), "expired"), + (_Verifier(result=AssertionRejection(reason="iss is not an issuer this gateway trusts")), "iss"), + (_Verifier(raises=Exception("Validation fails: Invalid audience")), "Invalid audience"), + (_Verifier(result=AssertionRejection(reason="bad kid\nWARNING forged line")), "bad kid WARNING forged line"), + ], +) +async def test_an_assertion_the_verifier_rejects_is_refused_with_the_reason_kept_in_the_log(verifier, reason, caplog): + resolver = _Resolver() + with caplog.at_level(logging.WARNING, logger="litellm.proxy"): + refused = await _identity(_redemption(_assertion()), verifier, resolver) + assert refused == SubjectTokenRefusal(error="invalid_grant", description=REJECTED_ASSERTION) + assert reason not in refused.description + assert reason in caplog.text + assert resolver.calls == () + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "claims, reason", + [ + ({"client_id": "0oa-confidential"}, "client_id"), + ({"exp": NOW}, "expired"), + ({"exp": NOW - 1}, "expired"), + ({"jti": ""}, "jti"), + ({"jti": 7}, "jti"), + ({"iss": None}, "iss"), + ({"cnf": {"jkt": "thumbprint"}}, "proof"), + ], +) +async def test_verified_claims_that_do_not_fit_the_profile_are_refused(claims, reason, caplog): + resolver = _Resolver() + with caplog.at_level(logging.WARNING, logger="litellm.proxy"): + refused = await _identity(_redemption(_assertion()), _Verifier(result=_claims(**claims)), resolver) + assert refused == SubjectTokenRefusal(error="invalid_grant", description=REJECTED_ASSERTION) + assert reason in caplog.text + assert resolver.calls == () + + +@pytest.mark.asyncio +async def test_an_assertion_that_names_no_user_the_gateway_knows_cannot_be_redeemed(caplog): + with caplog.at_level(logging.WARNING, logger="litellm.proxy"): + refused = await _identity(_redemption(_assertion()), resolver=_Resolver(user_id=None)) + assert refused == SubjectTokenRefusal(error="invalid_grant", description=REJECTED_ASSERTION) + assert "names no user" in caplog.text + missing = await _identity( + _redemption(_assertion()), resolver=_Resolver(raises=ValueError("User doesn't exist in db. 'user_id'=u1")) + ) + assert missing == SubjectTokenRefusal(error="invalid_grant", description=REJECTED_ASSERTION) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "verifier, resolver", + [ + (_Verifier(raises=JWKS_DOWN), _Resolver()), + (_Verifier(), _Resolver(raises=HTTPException(status_code=503, detail="database unavailable"))), + ], +) +async def test_a_gateway_outage_is_reported_as_retryable_not_as_a_bad_assertion(verifier, resolver): + refused = await _identity(_redemption(_assertion()), verifier, resolver) + assert isinstance(refused, SubjectTokenRefusal) + assert refused.error == "temporarily_unavailable" + assert "retry" in refused.description + assert "assertion" in refused.description + assert "subject_token" not in refused.description + + +@pytest.mark.parametrize( + "unmet, description", + [ + ({"jwt_auth_enabled": False}, "JWT auth is not enabled"), + ({"has_database": False}, "no database"), + ({"licensed": False}, "enterprise only"), + ({"maps_jwts_to_virtual_keys": True}, "maps IdP tokens to virtual keys"), + ({"shared_replay_store": False}, "Redis"), + ({"base_url": None}, "PROXY_BASE_URL"), + ({"requesting_apps": ()}, "no requesting apps allowlisted"), + ], +) +@pytest.mark.asyncio +async def test_the_deployment_gates_refuse_before_any_verification(unmet, description): + verifier = _Verifier() + refused = await _identity(_redemption(_assertion()), verifier, **unmet) + assert isinstance(refused, SubjectTokenRefusal) + assert refused.error == "unsupported_grant_type" + assert description in refused.description + assert verifier.calls == () + assert IdJagPrerequisites(**{**EVERY_GATE_HOLDS, **unmet}).available is False + assert IdJagPrerequisites(**EVERY_GATE_HOLDS).available is True + + +def _running_jwt_handler(litellm_jwtauth: LiteLLM_JWTAuth | None) -> JWTHandler: + handler: Final = JWTHandler() + if litellm_jwtauth is not None: + handler.update_environment(prisma_client=None, user_api_key_cache=DualCache(), litellm_jwtauth=litellm_jwtauth) + return handler + + +@pytest.mark.parametrize( + "general_settings, prisma_client, premium_user, litellm_jwtauth, expected", + [ + ({"enable_jwt_auth": True}, object(), True, LiteLLM_JWTAuth(id_jag_requesting_apps=[PUBLIC_APP]), True), + ( + {"enable_jwt_auth": True, "redis": None}, + object(), + True, + LiteLLM_JWTAuth(id_jag_requesting_apps=[PUBLIC_APP]), + False, + ), + ( + {"enable_jwt_auth": True, "base_url": ""}, + object(), + True, + LiteLLM_JWTAuth(id_jag_requesting_apps=[PUBLIC_APP]), + False, + ), + ({"enable_jwt_auth": True}, object(), True, LiteLLM_JWTAuth(id_jag_requesting_apps=[]), False), + ( + {"enable_jwt_auth": True}, + object(), + True, + LiteLLM_JWTAuth(id_jag_requesting_apps=[PUBLIC_APP], virtual_key_claim_field="vkey"), + False, + ), + ({"enable_jwt_auth": True}, object(), True, LiteLLM_JWTAuth(), False), + ({"enable_jwt_auth": True}, object(), True, None, False), + ({}, object(), True, LiteLLM_JWTAuth(id_jag_requesting_apps=[PUBLIC_APP]), False), + ({"enable_jwt_auth": True}, None, True, LiteLLM_JWTAuth(id_jag_requesting_apps=[PUBLIC_APP]), False), + ({"enable_jwt_auth": True}, object(), False, LiteLLM_JWTAuth(id_jag_requesting_apps=[PUBLIC_APP]), False), + ], +) +def test_availability_is_read_from_the_running_proxy( + monkeypatch, general_settings, prisma_client, premium_user, litellm_jwtauth, expected +): + monkeypatch.setenv("PROXY_BASE_URL", general_settings.get("base_url", "https://gateway.example.com")) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", premium_user) + monkeypatch.setattr("litellm.proxy.proxy_server.redis_usage_cache", general_settings.get("redis", object())) + monkeypatch.setattr("litellm.proxy.proxy_server.jwt_handler", _running_jwt_handler(litellm_jwtauth)) + assert jwt_bearer_available() is expected + + +def test_the_assertion_audience_is_the_fixed_gateway_url_not_the_request_host(): + """An ID-JAG must name this gateway's MCP resource or its authorization server under the + operator-fixed PROXY_BASE_URL. Without it there is no audience, so there is no grant.""" + assert gateway_audiences("https://gateway.example.com") == ( + "https://gateway.example.com/mcp", + "https://gateway.example.com", + ) + assert gateway_audiences(None) == () + + +@pytest.mark.asyncio +async def test_a_team_the_subject_lookup_cannot_settle_is_refused_before_the_assertion_is_consumed(): + """The resolver's refusal (a user on several teams who named none, or a header naming a team + they are not on) comes back as is, so the client fixes its request and presents the same + assertion again instead of finding its jti burned.""" + refusal = SubjectTokenRefusal(error="invalid_request", description=f"send the {TEAM_HEADER} header") + refused = await _identity(_redemption(_assertion()), resolver=_Resolver(result=refusal)) + assert refused is refusal + teamed = await _identity( + _redemption(_assertion()), resolver=_Resolver(result=SubjectIdentity(user_id="u1", team_id="t9")) + ) + assert isinstance(teamed, AssertionIdentity) + assert teamed.team_id == "t9" + + +ENG = ConsentTeam(team_id="team-eng", team_alias="engineering") +OPS = ConsentTeam(team_id="team-ops", team_alias=None) + + +@pytest.mark.parametrize( + "team_header, claim_teams, teams, expected", + [ + (None, (), (), None), + (None, (), (ENG,), "team-eng"), + ("team-eng", (), (ENG, OPS), "team-eng"), + ("engineering", (), (ENG, OPS), "team-eng"), + ("team-ops", ("team-eng",), (ENG, OPS), "team-ops"), + (None, ("team-ops",), (ENG, OPS), "team-ops"), + (None, ("team-sales", "team-ops"), (ENG, OPS), "team-ops"), + (None, ("0oa-xaa-public-app",), (ENG,), "team-eng"), + (None, ("0oa-xaa-public-app",), (), None), + ], +) +def test_select_team_prefers_the_header_then_the_assertion_then_the_only_team( + team_header, claim_teams, teams, expected +): + assert select_team(team_header, claim_teams, teams) == expected + + +class _Reader: + def __init__(self, user_id=None, email=None, team_ids=(), allowed_domain=None): + self.user_id, self.email, self.team_ids, self.allowed_domain = user_id, email, team_ids, allowed_domain + + def is_allowed_domain(self, user_email): + return self.allowed_domain is None or user_email.split("@")[-1] == self.allowed_domain + + def get_user_id(self, token, default_value): + return self.user_id or default_value + + def get_user_email(self, token, default_value): + return self.email or default_value + + def get_all_jwt_team_ids(self, token): + return list(self.team_ids) + + +async def _known_user(user_id: str, user_email: str | None) -> str | None: + return user_id if user_id in {"okta-user-1", "u-mapped"} else None + + +async def _two_teams(user_id: str) -> tuple[ConsentTeam, ...]: + return (ENG, OPS) + + +async def _teams_unreadable(user_id: str) -> ReloadUserFailure: + return "unavailable" + + +def _teams_failing(failure: ReloadUserFailure) -> LookupTeams: + async def lookup_teams(user_id: str) -> tuple[ConsentTeam, ...] | ReloadUserFailure: + return failure + + return lookup_teams + + +@pytest.mark.asyncio +async def test_the_subject_resolver_binds_the_claims_to_a_known_user_and_their_asserted_team(): + """The user comes from the JWT auth user mapping, else ``sub``, and must exist. The team the + assertion asserts settles a multi-team user, the header still wins, and unreadable teams are a + retryable refusal rather than a credential for the wrong team.""" + asserted = subject_resolver(_Reader(team_ids=("team-ops",)), _known_user, _two_teams) + assert await asserted(_claims(), None) == SubjectIdentity(user_id="okta-user-1", team_id="team-ops") + assert await asserted(_claims(), "team-eng") == SubjectIdentity(user_id="okta-user-1", team_id="team-eng") + mapped = subject_resolver(_Reader(user_id="u-mapped", email="m@example.com"), _known_user, _two_teams) + assert await mapped(_claims(), "engineering") == SubjectIdentity(user_id="u-mapped", team_id="team-eng") + unknown = subject_resolver(_Reader(), _known_user, _two_teams) + assert await unknown(_claims(sub="nobody"), None) is None + unreadable = await subject_resolver(_Reader(), _known_user, _teams_unreadable)(_claims(), None) + assert isinstance(unreadable, SubjectTokenRefusal) + assert unreadable.error == "temporarily_unavailable" + + +@pytest.mark.parametrize("failure", ["no_active_key", "unresolvable"]) +@pytest.mark.asyncio +async def test_a_blocked_or_unmapped_user_is_a_bad_grant_not_an_outage_to_retry(failure): + """A deactivated user, or one the resolver cannot map, can retry forever and never succeed, so + the refusal is invalid_grant. Only a database the gateway could not reach is retryable.""" + refused = await subject_resolver(_Reader(), _known_user, _teams_failing(failure))(_claims(), None) + assert isinstance(refused, SubjectTokenRefusal) + assert refused.error == "invalid_grant" + assert "retry" not in refused.description + outage = await subject_resolver(_Reader(), _known_user, _teams_failing("faulted"))(_claims(), None) + assert isinstance(outage, SubjectTokenRefusal) + assert outage.error == "temporarily_unavailable" + + +def test_the_assertion_team_claims_ignore_team_id_default(): + """``team_id_default`` is how bearer JWT auth picks a team for a request, not a claim the + assertion made, so a user outside that team must still get their only team or no team.""" + handler = JWTHandler() + handler.litellm_jwtauth = LiteLLM_JWTAuth(team_id_default="default-team") + claim_teams = tuple(_claim_reader(handler).get_all_jwt_team_ids(token={"iss": ISSUER, "sub": "u1"})) + assert claim_teams == () + assert select_team(None, claim_teams, (ENG,)) == "team-eng" + assert select_team(None, claim_teams, ()) is None + + +@pytest.mark.parametrize( + "team_header, claim_teams, teams, error, mention", + [ + ("team-sales", (), (ENG, OPS), "invalid_request", TEAM_HEADER), + ("engineering", (), (), "invalid_request", TEAM_HEADER), + (None, (), (ENG, OPS), "invalid_request", TEAM_HEADER), + (None, ("team-eng", "team-ops"), (ENG, OPS), "invalid_request", TEAM_HEADER), + (None, ("0oa-xaa-public-app",), (ENG, OPS), "invalid_request", TEAM_HEADER), + ], +) +def test_select_team_refuses_a_team_the_user_is_not_on_or_an_unnamed_choice( + team_header, claim_teams, teams, error, mention +): + refused = select_team(team_header, claim_teams, teams) + assert isinstance(refused, SubjectTokenRefusal) + assert refused.error == error + assert mention in refused.description + + +@pytest.mark.parametrize( + "authorization, expected", + [ + (None, None), + ("Bearer abc", None), + ("Basic not-base64!", None), + (_basic("0oa-confidential", "s3cret"), "0oa-confidential"), + ("basic " + b64encode(b"0oa%20x:s3cret").decode(), "0oa x"), + ("Basic " + b64encode(b"0oa+x:s3cret").decode(), "0oa x"), + ("Basic " + b64encode(b"no-separator").decode(), None), + ], +) +def test_basic_client_id_reads_the_app_from_the_authorization_header(authorization, expected): + assert basic_client_id(authorization) == expected + + +MCP_ROUTE_GROUPS = ("mcp_inference_routes",) +LLM_ROUTE_GROUPS = ("llm_api_routes_without_mcp",) + + +@pytest.mark.parametrize( + "claim, scope, allowed_routes", + [ + (None, None, None), + ("", None, None), + ("litellm:mcp", "litellm:mcp", MCP_ROUTE_GROUPS), + ("litellm:llm", "litellm:llm", LLM_ROUTE_GROUPS), + ("litellm:mcp litellm:llm", "litellm:mcp litellm:llm", MCP_ROUTE_GROUPS + LLM_ROUTE_GROUPS), + ("litellm:llm litellm:mcp litellm:llm", "litellm:mcp litellm:llm", MCP_ROUTE_GROUPS + LLM_ROUTE_GROUPS), + ], +) +def test_the_scope_claim_buys_exactly_the_route_groups_it_names(claim, scope, allowed_routes): + """``litellm:mcp`` is the MCP gateway's data-plane routes, ``litellm:llm`` the LLM API routes + without their MCP members, both together are both, and no claim is the full credential. The + scope echoed back is canonical: deduplicated, in vocabulary order, so it is exactly what is + enforced. Every group named must be one the route check resolves.""" + granted = granted_scope(claim, mcp_resource=False) + assert (granted.scope, granted.allowed_routes) == (scope, allowed_routes) + if allowed_routes is not None: + assert set(allowed_routes) <= set(LiteLLMRoutes._member_names_) + mcp_members = set(LiteLLMRoutes.mcp_inference_routes.value) + llm_members = set().union(*(LiteLLMRoutes[group].value for group in LLM_ROUTE_GROUPS)) + assert "/chat/completions" in llm_members and "/mcp" not in llm_members + assert "/mcp" in mcp_members and "/chat/completions" not in mcp_members + + +@pytest.mark.parametrize("claim", ["openid", "litellm:mcp openid", "litellm:admin", 7, ["litellm:mcp"]]) +@pytest.mark.asyncio +async def test_a_scope_the_gateway_does_not_know_is_refused_as_invalid_scope_before_any_lookup(claim): + refusal = granted_scope(claim, mcp_resource=False) + assert isinstance(refusal, SubjectTokenRefusal) + assert refusal.error == "invalid_scope" + assert "litellm:mcp" in refusal.description and "litellm:llm" in refusal.description + resolver = _Resolver() + refused = await _identity(_redemption(_assertion()), _Verifier(result=_claims(scope=claim)), resolver) + assert refused == refusal + assert resolver.calls == () + + +@pytest.mark.asyncio +async def test_a_linked_sso_account_resolves_by_its_provider_subject_without_an_email_claim(): + """An SSO login stores the provider subject as ``sso_user_id`` on a user whose ``user_id`` is + something else, so the lookup must offer the subject as both, with no email to fall back on.""" + from litellm.models.user import LiteLLM_UserTable + + async def get_user(*, user_id, user_email, sso_user_id): + if sso_user_id == "00u-okta-subject" and user_email is None: + return LiteLLM_UserTable(user_id="local-user-7", user_email=None) + return None + + lookup = proxy_user_lookup(get_user) + assert await lookup("00u-okta-subject", None) == "local-user-7" + assert await lookup("nobody", None) is None + + +@pytest.mark.parametrize("email", ["eve@evil.example", None]) +@pytest.mark.asyncio +async def test_the_allowed_email_domain_gates_the_assertions_subject_like_bearer_jwt_auth(email): + """With ``user_allowed_email_domain`` set, bearer JWT auth refuses an identity whose email is + outside the domain or missing. An assertion for an existing user gets the same refusal before + any lookup, as invalid_grant, and a matching email still resolves.""" + refused = await subject_resolver(_Reader(email=email, allowed_domain="corp.example"), _known_user, _two_teams)( + _claims(), "team-eng" + ) + assert isinstance(refused, SubjectTokenRefusal) + assert refused.error == "invalid_grant" + allowed = subject_resolver( + _Reader(email="ann@corp.example", allowed_domain="corp.example"), _known_user, _two_teams + ) + assert await allowed(_claims(), "team-eng") == SubjectIdentity(user_id="okta-user-1", team_id="team-eng") + + +@pytest.mark.parametrize("claim", [None, "", "litellm:mcp", "litellm:mcp litellm:mcp"]) +def test_an_assertion_for_the_mcp_resource_is_ceilinged_to_the_mcp_routes(claim): + """``aud = {base}/mcp`` means the identity provider named the MCP resource, so no scope or + ``litellm:mcp`` both mint an MCP-only credential and echo ``litellm:mcp``.""" + granted = granted_scope(claim, mcp_resource=True) + assert granted == GrantedScope(scope="litellm:mcp", allowed_routes=("mcp_inference_routes",)) + + +@pytest.mark.parametrize("claim", ["litellm:llm", "litellm:mcp litellm:llm"]) +def test_llm_routes_cannot_be_granted_under_the_mcp_resource_audience(claim): + refusal = granted_scope(claim, mcp_resource=True) + assert isinstance(refusal, SubjectTokenRefusal) + assert refusal.error == "invalid_scope" + assert "/mcp audience" in refusal.description and "bare origin" in refusal.description + + +@pytest.mark.parametrize( + "aud, expected", + [ + (AUDIENCES[0], True), + (AUDIENCES[1], False), + ([AUDIENCES[0]], True), + ([AUDIENCES[0], AUDIENCES[1]], False), + ("https://other.example", False), + (None, False), + ], +) +def test_the_mcp_resource_audience_is_the_mcp_path_without_the_bare_origin(aud, expected): + assert names_mcp_resource(aud, GATEWAY) is expected + assert names_mcp_resource(aud, None) is False + + +@pytest.mark.asyncio +async def test_redeeming_for_the_mcp_resource_narrows_the_identity_and_refuses_llm_scope(): + """End to end through identity_from_id_jag: the same verified claims, with ``aud`` the MCP + resource, come out ceilinged, and asking for LLM routes under that audience is a scope error.""" + narrowed = await _identity(_redemption(_assertion(aud=AUDIENCES[0])), _Verifier(), _Resolver()) + assert isinstance(narrowed, AssertionIdentity) + assert (narrowed.scope, narrowed.allowed_routes) == ("litellm:mcp", ("mcp_inference_routes",)) + refused = await _identity(_redemption(_assertion(aud=AUDIENCES[0], scope="litellm:llm")), _Verifier(), _Resolver()) + assert isinstance(refused, SubjectTokenRefusal) + assert refused.error == "invalid_scope" + unscoped = jwt.encode( + {key: value for key, value in _claims(aud=AUDIENCES[1]).items() if key != "scope"}, + "unused-by-the-fake-verifier-but-long-enough-for-hs256", + algorithm="HS256", + headers={"typ": ID_JAG_TYP}, + ) + full = await _identity(_redemption(unscoped), _Verifier(), _Resolver()) + assert isinstance(full, AssertionIdentity) + assert (full.scope, full.allowed_routes) == (None, None) + scoped = await _identity(_redemption(_assertion(aud=AUDIENCES[1], scope="litellm:llm")), _Verifier(), _Resolver()) + assert isinstance(scoped, AssertionIdentity) + assert scoped.scope == "litellm:llm" diff --git a/tests/unit/proxy/_experimental/mcp_server/test_proxy_api_credentials.py b/tests/unit/proxy/_experimental/mcp_server/test_proxy_api_credentials.py index 8484dfdde72..f39aec82c3e 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_proxy_api_credentials.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_proxy_api_credentials.py @@ -1,15 +1,21 @@ """Tests for minting the ``lite login`` credential from a consented native-client grant.""" -from unittest.mock import ANY, AsyncMock, MagicMock +from unittest.mock import ANY, AsyncMock, MagicMock, patch import pytest +from fastapi import HTTPException from litellm.constants import CLI_JWT_EXPIRATION_HOURS from litellm.models.user import LiteLLM_UserTable from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ConsentTeam, MintedProxyCredential -from litellm.proxy._experimental.mcp_server.proxy_api_credentials import lookup_consent_teams, mint_proxy_credential +from litellm.proxy._experimental.mcp_server.proxy_api_credentials import ( + lookup_consent_teams, + lookup_consent_teams_from_database, + mint_proxy_credential, +) from litellm.proxy._types import LitellmUserRoles from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken +from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.management_endpoints.ui_sso import CliSsoTeamDetail @@ -235,3 +241,82 @@ async def test_lookup_consent_teams_passes_user_lookup_failures_through(failure, async def test_lookup_consent_teams_reports_unavailable_when_details_cannot_load(load_user, fetch_teams): fetch_teams.return_value = None assert await lookup_consent_teams("u1") == "unavailable" + + +def _route_status(minted: MintedProxyCredential, route: str, method: str | None = None) -> int: + request = None if method is None else MagicMock(method=method) + try: + RouteChecks.is_virtual_key_allowed_to_call_route(route=route, valid_token=_decoded(minted), request=request) + except HTTPException as refused: + return refused.status_code + return 200 + + +@pytest.mark.parametrize( + "allowed_routes, chat_completions, mcp", + [ + (None, True, True), + (("mcp_inference_routes",), False, True), + (("llm_api_routes_without_mcp",), True, False), + (("mcp_inference_routes", "llm_api_routes_without_mcp"), True, True), + ], + ids=["full credential", "mcp only", "llm only", "both"], +) +@pytest.mark.asyncio +async def test_mint_narrows_the_credential_to_the_allowed_routes_the_route_check_enforces( + load_user, fetch_teams, allowed_routes, chat_completions, mcp +): + """The allowed_routes a grant's scope resolved to ride inside the credential, so the same route + check a virtual key goes through on every request refuses /chat/completions for an MCP-only + credential and /mcp for an LLM-only one, and a credential minted without them is unchanged.""" + minted = await mint_proxy_credential("u1", "team-a", allowed_routes=allowed_routes) + assert isinstance(minted, MintedProxyCredential) + assert _decoded(minted).allowed_routes == list(allowed_routes or ()) + chat_status, mcp_status = (200 if chat_completions else 403), (200 if mcp else 403) + assert _route_status(minted, "/chat/completions") == chat_status + assert _route_status(minted, "/v1/chat/completions") == chat_status + assert _route_status(minted, "/mcp") == mcp_status + assert _route_status(minted, "/mcp/tools/call") == mcp_status + + +@pytest.mark.asyncio +async def test_lookup_consent_teams_from_database_reads_the_row_the_minter_will_recheck(load_user, fetch_teams): + """A grant consumed once a team is picked must pick it from the same membership the minter re-reads, + or a stale cache would burn the assertion on a request the caller could have fixed.""" + teams = await lookup_consent_teams_from_database("u1") + load_user.assert_awaited_once_with("u1", source="database") + assert teams == (ConsentTeam(team_id="team-a", team_alias="Team A"), ConsentTeam(team_id="team-b")) + + +@pytest.mark.parametrize( + "allowed_routes, registered_pass_through, auto_router_session, mcp_server_listing, agent_registry", + [ + (("llm_api_routes",), 200, 200, 200, 200), + (("llm_api_routes_without_mcp",), 200, 200, 403, 403), + (("mcp_inference_routes",), 403, 403, 403, 403), + ], + ids=["llm_api_routes key", "llm only", "mcp only"], +) +@pytest.mark.asyncio +async def test_an_llm_only_credential_keeps_the_llm_api_carve_outs_except_mcp_discovery( + load_user, + fetch_teams, + allowed_routes, + registered_pass_through, + auto_router_session, + mcp_server_listing, + agent_registry, +): + """Operator-registered pass-through endpoints and the auto-router session read are LLM API + surface that only the ``llm_api_routes`` group used to reach, so an LLM-only credential keeps + them, while the MCP server listing that group also carves out stays MCP surface.""" + minted = await mint_proxy_credential("u1", "team-a", allowed_routes=allowed_routes) + assert isinstance(minted, MintedProxyCredential) + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.InitPassThroughEndpointHelpers.is_registered_pass_through_route", + side_effect=lambda route: route.startswith("/my-llm-endpoint"), + ): + assert _route_status(minted, "/my-llm-endpoint/v1/chat", "POST") == registered_pass_through + assert _route_status(minted, "/auto_router/session", "GET") == auto_router_session + assert _route_status(minted, "/v1/mcp/server", "GET") == mcp_server_listing + assert _route_status(minted, "/v1/agents", "GET") == agent_registry diff --git a/tests/unit/proxy/auth/test_handle_jwt.py b/tests/unit/proxy/auth/test_handle_jwt.py index 69eb01cabcd..e4a3935a58d 100644 --- a/tests/unit/proxy/auth/test_handle_jwt.py +++ b/tests/unit/proxy/auth/test_handle_jwt.py @@ -13,6 +13,7 @@ from litellm.caching.dual_cache import DualCache from litellm.proxy._types import ( DEFAULT_JWKS_STALE_TTL, JWTLiteLLMRoleMap, + JWTIssuerConfig, LiteLLMRoutes, LiteLLM_JWTAuth, LiteLLM_ModelTable, @@ -29,9 +30,11 @@ from litellm.proxy._types import ( from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry from litellm.proxy.auth.auth_checks import TeamNotFoundError from litellm.proxy.auth.handle_jwt import ( + ID_JAG_TYP, JWKS_FETCH_ATTEMPTS, STALE_CACHE_KEY_PREFIX, STALE_WRITTEN_AT_CACHE_KEY_PREFIX, + AssertionRejection, HeaderTeam, JWKSUnreachableError, JWTAuthManager, @@ -8303,3 +8306,198 @@ def test_team_passthrough_route_denial_names_the_matched_deny_entry( assert exc_info.value.status_code == 403 assert expected_detail in str(exc_info.value.detail) + + +ASSERTION_ISSUER = "https://idp.example.com/oauth2/default" +ASSERTION_JWKS_URL = f"{ASSERTION_ISSUER}/v1/keys" +GATEWAY_AUDIENCES = ("https://gateway.example.com/mcp", "https://gateway.example.com") +ID_JAG_CLAIMS = ("iss", "sub", "aud", "exp", "iat", "jti", "client_id") + + +def _assertion_signer(): + import json + + import jwt as pyjwt + from cryptography.hazmat.primitives.asymmetric import rsa + + private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + jwk: Final = {**json.loads(pyjwt.algorithms.RSAAlgorithm.to_jwk(private_key.public_key())), "kid": "k1"} + + def sign(claims: Mapping[str, object], typ: str = ID_JAG_TYP) -> str: + return pyjwt.encode(dict(claims), private_key, algorithm="RS256", headers={"kid": "k1", "typ": typ}) + + return jwk, sign + + +def _id_jag_claims(**overrides): + now: Final = int(time.time()) + return { + "iss": ASSERTION_ISSUER, + "sub": "okta-user-1", + "aud": GATEWAY_AUDIENCES[0], + "exp": now + 300, + "iat": now - 5, + "jti": "jti-1", + "client_id": "0oa-requesting-app", + **overrides, + } + + +def _assertion_handler(jwk, litellm_jwtauth, jwks_url=ASSERTION_JWKS_URL): + """A handler whose JWKS for the issuer is already cached, so no HTTP is needed.""" + handler: Final = JWTHandler() + cache: Final = DualCache() + cache.set_cache(f"litellm_jwt_auth_keys_{jwks_url}", [jwk]) + handler.update_environment(prisma_client=None, user_api_key_cache=cache, litellm_jwtauth=litellm_jwtauth) + return handler + + +@pytest.mark.asyncio +async def test_decode_trusted_assertion_verifies_against_the_configured_issuer_with_the_callers_audience(): + """The assertion is checked with the issuer's JWKS and claim mappings, but for the audience the + caller names (the gateway's own issuer values), not the API audience the issuer entry carries, + since an ID-JAG is addressed to the authorization server, never to the API.""" + jwk, sign = _assertion_signer() + handler = _assertion_handler( + jwk, + LiteLLM_JWTAuth( + issuers=[ + JWTIssuerConfig( + issuer=ASSERTION_ISSUER, + jwks_url=ASSERTION_JWKS_URL, + audience="api://litellm", + user_id_jwt_field="sub", + ) + ] + ), + ) + claims = _id_jag_claims() + verified = await handler.decode_trusted_assertion(sign(claims), GATEWAY_AUDIENCES, ID_JAG_CLAIMS) + assert isinstance(verified, dict) + assert {key: verified[key] for key in claims} == claims + assert handler.get_user_id(token=verified, default_value=None) == "okta-user-1" + second = await handler.decode_trusted_assertion( + sign(_id_jag_claims(aud=GATEWAY_AUDIENCES[1])), GATEWAY_AUDIENCES, ID_JAG_CLAIMS + ) + assert isinstance(second, dict) + + +@pytest.mark.asyncio +async def test_decode_trusted_assertion_falls_back_to_the_env_scoped_issuer(monkeypatch): + jwk, sign = _assertion_signer() + monkeypatch.setenv("JWT_ISSUER", ASSERTION_ISSUER) + monkeypatch.setenv("JWT_PUBLIC_KEY_URL", ASSERTION_JWKS_URL) + handler = _assertion_handler(jwk, LiteLLM_JWTAuth()) + verified = await handler.decode_trusted_assertion(sign(_id_jag_claims()), GATEWAY_AUDIENCES, ID_JAG_CLAIMS) + assert isinstance(verified, dict) + assert verified["jti"] == "jti-1" + monkeypatch.setenv("JWT_ISSUER", "https://someone-else.example.com") + foreign = await handler.decode_trusted_assertion(sign(_id_jag_claims()), GATEWAY_AUDIENCES, ID_JAG_CLAIMS) + assert isinstance(foreign, AssertionRejection) + assert "trusts" in foreign.reason + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "fault, reason", + [ + ({"aud": "api://litellm"}, "udience"), + ({"jti": None}, "jti"), + ({"client_id": None}, "client_id"), + ({"exp": -10}, "xpired"), + ({"iat": 600}, "iat"), + ({"iss": "https://untrusted.example.com"}, "trusts"), + ], +) +async def test_decode_trusted_assertion_rejects_a_claim_outside_the_profile(fault: dict[str, object], reason: str): + """The claims are built inside the test, relative to the clock the verifier reads, so a slow + collection cannot turn an intended claim fault into an unrelated expiry (``None`` drops a claim, + an int on ``exp`` or ``iat`` is an offset from now).""" + now: Final = int(time.time()) + offsets: Final = {k: now + v for k, v in fault.items() if k in ("exp", "iat") and isinstance(v, int)} + claims: Final = {k: v for k, v in _id_jag_claims(**offsets).items() if fault.get(k, k) is not None} + claims.update({k: v for k, v in fault.items() if v is not None and k not in offsets}) + jwk, sign = _assertion_signer() + handler = _assertion_handler( + jwk, + LiteLLM_JWTAuth(issuers=[JWTIssuerConfig(issuer=ASSERTION_ISSUER, jwks_url=ASSERTION_JWKS_URL, audience="x")]), + ) + rejected = await handler.decode_trusted_assertion(sign(claims), GATEWAY_AUDIENCES, ID_JAG_CLAIMS) + assert isinstance(rejected, AssertionRejection) + assert reason in rejected.reason + + +@pytest.mark.asyncio +async def test_decode_trusted_assertion_rejects_a_forged_signature_and_garbage(monkeypatch): + jwk, _ = _assertion_signer() + _, forge = _assertion_signer() + handler = _assertion_handler( + jwk, + LiteLLM_JWTAuth(issuers=[JWTIssuerConfig(issuer=ASSERTION_ISSUER, jwks_url=ASSERTION_JWKS_URL, audience="x")]), + ) + forged = await handler.decode_trusted_assertion(forge(_id_jag_claims()), GATEWAY_AUDIENCES, ID_JAG_CLAIMS) + assert isinstance(forged, AssertionRejection) + assert "verification" in forged.reason + garbage = await handler.decode_trusted_assertion("not.a.jwt", GATEWAY_AUDIENCES, ID_JAG_CLAIMS) + assert isinstance(garbage, AssertionRejection) + monkeypatch.delenv("JWT_ISSUER", raising=False) + unconfigured = await _assertion_handler(jwk, LiteLLM_JWTAuth()).decode_trusted_assertion( + forge(_id_jag_claims()), GATEWAY_AUDIENCES, ID_JAG_CLAIMS + ) + assert isinstance(unconfigured, AssertionRejection) + assert "trusts" in unconfigured.reason + + +@pytest.mark.asyncio +async def test_decode_trusted_assertion_reports_an_unreachable_jwks_as_an_outage_not_a_bad_assertion(): + jwk, sign = _assertion_signer() + handler = _assertion_handler( + jwk, + LiteLLM_JWTAuth(issuers=[JWTIssuerConfig(issuer=ASSERTION_ISSUER, jwks_url=ASSERTION_JWKS_URL, audience="x")]), + jwks_url="https://a-different-cached-url.example.com/keys", + ) + with patch.object(handler, "_fetch_jwks_keys", new_callable=AsyncMock, side_effect=JWKSUnreachableError("down")): + with pytest.raises(ProxyException) as outage: + await handler.decode_trusted_assertion(sign(_id_jag_claims()), GATEWAY_AUDIENCES, ID_JAG_CLAIMS) + assert outage.value.code == "503" + + +@pytest.mark.asyncio +async def test_auth_jwt_never_accepts_an_id_jag_as_a_bearer_token(): + """An assertion typed oauth-id-jag+jwt is a grant for the token endpoint, so bearer auth + refuses it before any key lookup, even where it would otherwise verify.""" + jwk, sign = _assertion_signer() + handler = _assertion_handler( + jwk, + LiteLLM_JWTAuth( + issuers=[ + JWTIssuerConfig(issuer=ASSERTION_ISSUER, jwks_url=ASSERTION_JWKS_URL, audience=GATEWAY_AUDIENCES[0]) + ] + ), + ) + accepted = await handler.auth_jwt(sign(_id_jag_claims(), typ="JWT")) + assert accepted["jti"] == "jti-1" + with pytest.raises(Exception, match="token endpoint"): + await handler.auth_jwt(sign(_id_jag_claims())) +@pytest.mark.asyncio +@pytest.mark.parametrize("verdict", [True, False]) +async def test_decode_trusted_assertion_runs_the_configured_custom_validate_on_the_verified_claims(verdict): + """The operator's ``custom_validate`` gates every JWT the proxy admits, so an assertion redeemed + at the token endpoint goes through it too, after signature and claim verification.""" + jwk, sign = _assertion_signer() + seen: Final = [] + handler = _assertion_handler( + jwk, + LiteLLM_JWTAuth( + issuers=[JWTIssuerConfig(issuer=ASSERTION_ISSUER, jwks_url=ASSERTION_JWKS_URL, audience="api://litellm")] + ), + ) + handler.litellm_jwtauth.custom_validate = lambda claims: seen.append(dict(claims)) or verdict + claims = _id_jag_claims() + verified = await handler.decode_trusted_assertion(sign(claims), GATEWAY_AUDIENCES, ID_JAG_CLAIMS) + assert [{key: call[key] for key in claims} for call in seen] == [claims] + if verdict: + assert isinstance(verified, dict) + else: + assert isinstance(verified, AssertionRejection) + assert "custom_validate" in verified.reason diff --git a/tests/unit/proxy/test__lazy_features.py b/tests/unit/proxy/test__lazy_features.py index 4e0889fd93d..adc4de8c8fa 100644 --- a/tests/unit/proxy/test__lazy_features.py +++ b/tests/unit/proxy/test__lazy_features.py @@ -401,3 +401,25 @@ def test_without_the_flag_the_warmup_route_registers_a_feature_and_returns_its_p assert body.paths["/alpha/list"]["get"].tags == ("alpha",) assert loaded_lazy_modules(app) == {features[0].module_path} assert "/alpha/list" in _paths(app) and "/beta/list" not in _paths(app) + + +def test_root_authorization_server_metadata_is_the_gateways_not_the_byok_connector_document( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Two features claim the root ``/.well-known/oauth-authorization-server``. The one a client + gets must be the gateway's own authorization server metadata, whose ``/token`` serves the + gateway's grants, not the BYOK connector form flow whose token endpoint is + ``/v1/mcp/oauth/token``; the connector keeps its path-appended document.""" + monkeypatch.setenv(FLAG, "false") + app: Final = FastAPI() + attach_lazy_features(app) + + with TestClient(app, base_url="https://gateway.example.com") as client: + root: Final = client.get("/.well-known/oauth-authorization-server") + connector: Final = client.get("/.well-known/oauth-authorization-server/v1/mcp/oauth") + + assert root.status_code == 200 + assert root.json()["issuer"] == "https://gateway.example.com" + assert root.json()["token_endpoint"] == "https://gateway.example.com/token" + assert connector.status_code == 200 + assert connector.json()["token_endpoint"] == "https://gateway.example.com/v1/mcp/oauth/token" diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 6f6eae1f98e..05a4893ad7f 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -27480,8 +27480,10 @@ export interface components { }; /** Body_token_endpoint__mcp_server_name__token_post */ Body_token_endpoint__mcp_server_name__token_post: { + /** Assertion */ + assertion?: string | null; /** Client Id */ - client_id: string; + client_id?: string | null; /** Client Secret */ client_secret?: string | null; /** Code */ @@ -27507,8 +27509,10 @@ export interface components { }; /** Body_token_endpoint_token_post */ Body_token_endpoint_token_post: { + /** Assertion */ + assertion?: string | null; /** Client Id */ - client_id: string; + client_id?: string | null; /** Client Secret */ client_secret?: string | null; /** Code */ From 3501e36fa95d6e13aa5b70d6bd9d8e0a4d6f6c57 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Fri, 16 Oct 2026 12:16:56 -0700 Subject: [PATCH 13/58] feat(sso): capture the IdP id_token on Microsoft and Google logins for oauth2_id_jag (#46747) Resolves LIT-9671 Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../_experimental/mcp_server/oauth_utils.py | 5 +- .../sso_assertion_refresher.py | 93 ++++--- .../sso_assertion_store.py | 8 + .../sso/id_jag_assertion_capture.py | 25 +- litellm/proxy/management_endpoints/ui_sso.py | 206 +++++++++------ .../test_sso_assertion_refresher.py | 61 +++++ .../mcp_server/test_mcp_server_manager.py | 25 +- .../test_id_jag_assertion_capture.py | 50 ++-- .../test_mcp_management_endpoints.py | 17 +- .../proxy/management_endpoints/test_ui_sso.py | 241 +++++++++++++++--- 10 files changed, 537 insertions(+), 194 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/oauth_utils.py b/litellm/proxy/_experimental/mcp_server/oauth_utils.py index 69732186d75..133bad780bf 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_utils.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_utils.py @@ -19,6 +19,9 @@ from litellm.proxy._experimental.mcp_server.auth.token_endpoint_auth import ( build_token_endpoint_client_auth, normalize_token_endpoint_auth_method, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( + GOOGLE_REFRESH_TOKEN_PARAMS, +) from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.proxy.middleware.per_request_root_path_middleware import get_request_root_path @@ -59,7 +62,7 @@ _DEFAULT_NATIVE_REDIRECT_URIS: Final[list[str]] = [ ] _GOOGLE_AUTHORIZATION_HOSTS: Final = frozenset({"accounts.google.com"}) -_GOOGLE_OFFLINE_ACCESS_PARAMS: Final = MappingProxyType({"access_type": "offline", "prompt": "consent"}) +_GOOGLE_OFFLINE_ACCESS_PARAMS: Final = GOOGLE_REFRESH_TOKEN_PARAMS _warned_invalid_proxy_base_url: str | None = None diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_refresher.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_refresher.py index bca5848febe..c41730ce9ec 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_refresher.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_refresher.py @@ -12,15 +12,18 @@ not the size of the user table) and single-flighted per user through the same `` the ``authorization_code`` arm uses, because an IdP that rotates refresh tokens treats two concurrent redemptions of one token as replay and can revoke the whole grant chain. -The refresh is redeemed against the generic-OIDC client the login itself used -(``GENERIC_TOKEN_ENDPOINT`` / ``GENERIC_CLIENT_ID`` / ``GENERIC_CLIENT_SECRET``, which the proxy -reconciles from the stored SSO row into the process environment at startup), authenticated the way -that login authenticated: the non-PKCE path always sends HTTP Basic, while the PKCE path sends the -credentials in the body when ``GENERIC_INCLUDE_CLIENT_ID`` is set, and an IdP application may accept -only one of the two. An assertion can only exist if that client minted it, so no other client could -redeem its refresh token, and no other method is known to be accepted. A deployment whose -``GENERIC_SCOPE`` omits ``offline_access`` captures no refresh token at all, which is why that miss -logs the scope by name rather than failing silently. +The refresh is redeemed against the SSO client the login itself used, resolved in the same order +the login callback dispatches (Google, then Microsoft, then generic OIDC, from the client ids the +proxy reconciles from the stored SSO row into the process environment at startup), authenticated the +way that login authenticated: the Google and Microsoft logins and the non-PKCE generic path send +HTTP Basic, while the PKCE generic path sends the credentials in the body when +``GENERIC_INCLUDE_CLIENT_ID`` is set, and an IdP application may accept only one of the two. An +assertion can only exist if that client minted it, so no other client could redeem its refresh +token, and no other method is known to be accepted. A login whose authorization request never asked +for a refresh token captures none at all: the Google and Microsoft logins ask (``access_type=offline``, +``offline_access``) only while an ``oauth2_id_jag`` server exists, and generic OIDC asks only when +``GENERIC_SCOPE`` carries ``offline_access``, which is why that miss names the cause rather than +failing silently. Failures are values internally (``Result[_, RefreshFailure]``). At the store boundary they collapse onto the protocol's existing two-outcome contract: a refusal returns the expired assertion unchanged @@ -92,6 +95,10 @@ _SINGLE_FLIGHT_KEY: Final = "sso_identity_assertion" # Renew this far ahead of ``exp`` so a token that would die between resolution and the second leg of # the exchange is replaced first. Matches the sibling per-user token store's skew. _DEFAULT_EXPIRY_SKEW_SECONDS: Final = 60.0 +# Where fastapi-sso's Google login redeemed the authorization code, read from Google's discovery +# document; Microsoft's mirrors ``CustomMicrosoftSSO``, override first, tenant template second. +_GOOGLE_TOKEN_ENDPOINT: Final = "https://oauth2.googleapis.com/token" +_MICROSOFT_TOKEN_ENDPOINT_TEMPLATE: Final = "https://login.microsoftonline.com/{tenant}/oauth2/v2.0/token" class AssertionRead(Protocol): @@ -122,7 +129,7 @@ class FormPost(Protocol): @dataclass(frozen=True, slots=True) class SSOClientConfig: - """The generic-OIDC client credentials a refresh_token grant has to authenticate as, and how.""" + """The SSO client credentials a refresh_token grant has to authenticate as, and how.""" token_endpoint: str client_id: str @@ -131,32 +138,58 @@ class SSOClientConfig: def sso_client_config(env: Mapping[str, str]) -> SSOClientConfig | None: - """The configured generic-OIDC client, or ``None`` when the deployment has none. + """The configured SSO client the login used, or ``None`` when the deployment has none. - Read from the process environment because that is where the login path reads it - (``_setup_generic_sso_env_vars``) and where the proxy materializes the stored ``sso_config`` row - at startup, so this resolves to the same client that minted the assertion. ``None`` is an - ordinary state, not an error: a deployment signing in through a provider that captures no + Read from the process environment because that is where the login callback reads it and where + the proxy materializes the stored ``sso_config`` row at startup, and resolved in the callback's + own order (``GOOGLE_CLIENT_ID``, then ``MICROSOFT_CLIENT_ID``, then ``GENERIC_CLIENT_ID``, each + selecting its branch when set at all), so this is the same client that minted the assertion: a + deployment carrying both a Google and a generic client signed the user in through Google. ``None`` + is an ordinary state, not an error: a deployment signing in through a provider that captures no assertion has nothing here to renew, and a client with no secret is not a confidential client that could redeem one. - ``auth_method`` is derived from the same ``GENERIC_INCLUDE_CLIENT_ID`` the login reads, because - the two login paths do not agree: the non-PKCE path always authenticates with HTTP Basic, while - the PKCE path puts the credentials in the body when that flag is set. Both capture assertions, so - a constant here would authenticate the renewal differently from the sign-in that produced the + ``auth_method`` follows how each login authenticated at the token endpoint. The Google and + Microsoft logins (fastapi-sso) and the non-PKCE generic path send HTTP Basic, while the PKCE + generic path puts the credentials in the body when ``GENERIC_INCLUDE_CLIENT_ID`` is set. A + constant here would authenticate the renewal differently from the sign-in that produced the refresh token and 401 against an IdP application registered for only one of the two. """ - token_endpoint: Final = env.get("GENERIC_TOKEN_ENDPOINT") - client_id: Final = env.get("GENERIC_CLIENT_ID") - client_secret: Final = env.get("GENERIC_CLIENT_SECRET") + if env.get("GOOGLE_CLIENT_ID") is not None: + return _confidential_client( + _GOOGLE_TOKEN_ENDPOINT, env.get("GOOGLE_CLIENT_ID"), env.get("GOOGLE_CLIENT_SECRET"), "client_secret_basic" + ) + if env.get("MICROSOFT_CLIENT_ID") is not None: + tenant: Final = env.get("MICROSOFT_TENANT") + return _confidential_client( + env.get("MICROSOFT_TOKEN_ENDPOINT") + or (_MICROSOFT_TOKEN_ENDPOINT_TEMPLATE.format(tenant=tenant) if tenant else None), + env.get("MICROSOFT_CLIENT_ID"), + env.get("MICROSOFT_CLIENT_SECRET"), + "client_secret_basic", + ) + includes_client_id: Final = env.get("GENERIC_INCLUDE_CLIENT_ID", "false").lower() == "true" + return _confidential_client( + env.get("GENERIC_TOKEN_ENDPOINT"), + env.get("GENERIC_CLIENT_ID"), + env.get("GENERIC_CLIENT_SECRET"), + "client_secret_post" if includes_client_id else "client_secret_basic", + ) + + +def _confidential_client( + token_endpoint: str | None, + client_id: str | None, + client_secret: str | None, + auth_method: MCPTokenEndpointAuthMethod, +) -> SSOClientConfig | None: if not token_endpoint or not client_id or not client_secret: return None - includes_client_id: Final = env.get("GENERIC_INCLUDE_CLIENT_ID", "false").lower() == "true" return SSOClientConfig( token_endpoint=token_endpoint, client_id=client_id, client_secret=SecretStr(client_secret), - auth_method="client_secret_post" if includes_client_id else "client_secret_basic", + auth_method=auth_method, ) @@ -263,8 +296,9 @@ class SSOAssertionRefresher: if assertion.refresh_token is None: verbose_proxy_logger.warning( "ID-JAG: the stored IdP identity assertion for user_id=%s has expired and no refresh token was " - "captured with it, so it cannot be renewed without another interactive sign-in. Add " - "'offline_access' to GENERIC_SCOPE so the SSO login captures one.", + "captured with it, so it cannot be renewed without another interactive sign-in. Google and Microsoft " + "logins ask for one once an oauth2_id_jag MCP server is registered, so a user who signed in before " + "that has to sign in again. Generic OIDC asks only with 'offline_access' in GENERIC_SCOPE.", user_id, ) return Error(RefreshFailure.of_rejected("no refresh token was captured at sign-in")) @@ -272,11 +306,12 @@ class SSOAssertionRefresher: if config is None: verbose_proxy_logger.warning( "ID-JAG: the stored IdP identity assertion for user_id=%s has expired and cannot be renewed " - "because the generic SSO client is not configured (GENERIC_TOKEN_ENDPOINT, GENERIC_CLIENT_ID, " - "GENERIC_CLIENT_SECRET).", + "because the active SSO provider has no confidential client configured (GOOGLE_CLIENT_ID and " + "GOOGLE_CLIENT_SECRET, MICROSOFT_CLIENT_ID with MICROSOFT_CLIENT_SECRET and MICROSOFT_TENANT, or " + "GENERIC_TOKEN_ENDPOINT with GENERIC_CLIENT_ID and GENERIC_CLIENT_SECRET).", user_id, ) - return Error(RefreshFailure.of_rejected("the generic SSO client is not configured")) + return Error(RefreshFailure.of_rejected("the SSO client is not configured")) carried_refresh_token: Final = assertion.refresh_token.get_secret_value() # Whichever method the SSO login used for this client, since that is the one the IdP diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py index 0bc17d94165..ef30b3d75b6 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py @@ -20,6 +20,7 @@ from __future__ import annotations import json from collections.abc import Mapping, Sequence from datetime import datetime, timezone +from types import MappingProxyType from typing import TYPE_CHECKING, Final, Protocol import jwt @@ -35,6 +36,13 @@ if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient +"""What a Google or Microsoft authorization request has to carry for the IdP to issue a refresh token +beside the id_token. Asked for only while an ``oauth2_id_jag`` server exists to renew an assertion for, so +a deployment without one sees no extra consent.""" +MICROSOFT_REFRESH_TOKEN_SCOPE: Final = "offline_access" +GOOGLE_REFRESH_TOKEN_PARAMS: Final = MappingProxyType({"access_type": "offline", "prompt": "consent"}) + + _ASSERTION_DECRYPT_LOG_KEY: Final = "sso_identity_assertion" _STR_ADAPTER: Final[TypeAdapter[str]] = TypeAdapter(str) _MAYBE_STR_ADAPTER: Final[TypeAdapter[str | None]] = TypeAdapter(str | None) diff --git a/litellm/proxy/management_endpoints/sso/id_jag_assertion_capture.py b/litellm/proxy/management_endpoints/sso/id_jag_assertion_capture.py index 404fdfc83a9..deb9d2975de 100644 --- a/litellm/proxy/management_endpoints/sso/id_jag_assertion_capture.py +++ b/litellm/proxy/management_endpoints/sso/id_jag_assertion_capture.py @@ -1,25 +1,26 @@ """Whether the SSO provider the login callback dispatches to can capture an IdP identity assertion. An ``oauth2_id_jag`` MCP server spends the ``id_token`` captured at SSO login as its RFC 8693 -subject token. Only the generic OIDC login path reaches a token response the gateway retains one -from, so a deployment whose SSO runs through Google, Microsoft or SAML never stores an assertion -and every store-sourced ID-JAG exchange fails for every user, however many times they sign in. -Neither side can see that alone: the MCP registration knows nothing about SSO and the login knows -nothing about MCP. This module is the one shared answer both warn from. +subject token. The Google, Microsoft and generic OIDC login paths all reach a token response the +gateway retains one from; SAML has no id_token to capture, and a deployment with no SSO never signs +anyone in, so under either every store-sourced ID-JAG exchange fails for every user, however many +times they sign in. Neither side can see that alone: the MCP registration knows nothing about SSO +and the login knows nothing about MCP. This module is the one shared answer both warn from. """ from __future__ import annotations import os from enum import Enum +from typing import Final from typing_extensions import assert_never from litellm.proxy.management_endpoints.sso.saml_sso import SAMLAuthHandler -_GENERIC_OIDC_REMEDY = ( - "Point SSO at the generic OIDC provider (GENERIC_CLIENT_ID), the one login path whose token " - "response the gateway retains an id_token from" +_OIDC_REMEDY: Final = ( + "Point SSO at an OIDC provider (GOOGLE_CLIENT_ID, MICROSOFT_CLIENT_ID or GENERIC_CLIENT_ID), the login " + "paths whose token response the gateway retains an id_token from" ) @@ -55,18 +56,18 @@ def id_jag_assertion_capture_gap() -> str | None: or ``None`` when that provider does capture an assertion.""" provider = active_sso_provider() match provider: - case ActiveSSOProvider.generic: + case ActiveSSOProvider.google | ActiveSSOProvider.microsoft | ActiveSSOProvider.generic: return None case ActiveSSOProvider.none: return ( "no SSO provider is configured, so no IdP identity assertion is ever captured and " - f"ID-JAG credential resolution fails for every user. {_GENERIC_OIDC_REMEDY}" + f"ID-JAG credential resolution fails for every user. {_OIDC_REMEDY}" ) - case ActiveSSOProvider.google | ActiveSSOProvider.microsoft | ActiveSSOProvider.saml: + case ActiveSSOProvider.saml: return ( f"the active SSO provider ({provider.value}) has no identity-assertion capture path, so no " "IdP id_token is ever stored and ID-JAG credential resolution fails for every user no matter " - f"how often they sign in. {_GENERIC_OIDC_REMEDY}" + f"how often they sign in. {_OIDC_REMEDY}" ) case _: assert_never(provider) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 4629103070f..c3c5d7401c5 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -37,10 +37,12 @@ from urllib.parse import parse_qs, urlencode, urlparse if TYPE_CHECKING: import httpx + from litellm.caching.redis_cache import RedisCache + import jwt from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response, status from fastapi.responses import RedirectResponse -from pydantic import ConfigDict, TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError, field_validator import litellm from litellm._internal_context import with_service_target @@ -70,6 +72,8 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( + GOOGLE_REFRESH_TOKEN_PARAMS, + MICROSOFT_REFRESH_TOKEN_SCOPE, SSOIdentityAssertion, assertion_from_sso_login, ema_assertion_retention_enabled, @@ -1786,6 +1790,19 @@ async def warn_if_id_jag_assertion_uncaptured( ) +async def login_asks_for_refresh_token(retention_enabled: RetentionCheck | None = None) -> bool: + """A refresh token is only issued when the authorization request asks for one, and asking costs the + user an offline-access consent, so the Google and Microsoft logins ask only while an ``oauth2_id_jag`` + server exists whose assertion the refresher would renew. A store outage answers no: a login must not + fail over a renewal that is only an optimisation on top of signing in again.""" + try: + check: Final = retention_enabled if retention_enabled is not None else ema_assertion_retention_enabled + return await check() + except Exception as exc: # noqa: BLE001 # the login must not fail over the renewal gate + verbose_proxy_logger.debug("Could not check for oauth2_id_jag MCP servers before SSO login: %s", exc) + return False + + async def warn_if_id_jag_capture_gap(*, retention_enabled: RetentionCheck | None = None) -> None: gap: Final = id_jag_assertion_capture_gap() if gap is None: @@ -2125,6 +2142,50 @@ async def check_and_update_if_proxy_admin_id(user_role: str, user_id: str, prism return user_role +async def _sso_provider_callback( + request: Request, + redirect_url: str, + *, + google_client_id: str | None, + microsoft_client_id: str | None, + generic_client_id: str | None, + jwt_handler: JWTHandler, + sso_jwt_handler: JWTHandler | None, + return_raw_sso_response: bool = False, +) -> tuple[ + OpenID | dict[str, object] | None, dict[str, object] | None, dict[str, object] | None, SSOIdentityAssertion | None +]: + """Dispatch the callback to the configured provider in the order the login redirect used (Google, + Microsoft, generic) and hand back what every branch has to agree on: the parsed identity, the raw + userinfo and access-token claims only the generic branch exposes, and the IdP identity assertion + an ``oauth2_id_jag`` MCP server spends. ``None`` identity means no provider is configured.""" + if google_client_id is not None: + google_result, google_assertion = await GoogleSSOHandler.get_google_callback_response( + request=request, + google_client_id=google_client_id, + redirect_url=redirect_url, + return_raw_sso_response=return_raw_sso_response, + ) + return google_result, None, None, google_assertion + if microsoft_client_id is not None: + microsoft_result, microsoft_assertion = await MicrosoftSSOHandler.get_microsoft_callback_response( + request=request, + microsoft_client_id=microsoft_client_id, + redirect_url=redirect_url, + return_raw_sso_response=return_raw_sso_response, + ) + return microsoft_result, None, None, microsoft_assertion + if generic_client_id is not None: + return await get_generic_sso_response( + request=request, + jwt_handler=jwt_handler, + generic_client_id=generic_client_id, + redirect_url=redirect_url, + sso_jwt_handler=sso_jwt_handler, + ) + return None, None, None, None + + @router.get("/sso/callback", tags=["experimental"], include_in_schema=False) async def auth_callback(request: Request, state: str | None = None): """Verify login""" @@ -2173,9 +2234,6 @@ async def auth_callback(request: Request, state: str | None = None): microsoft_client_id: Final = os.getenv("MICROSOFT_CLIENT_ID", None) google_client_id: Final = os.getenv("GOOGLE_CLIENT_ID", None) generic_client_id: Final = os.getenv("GENERIC_CLIENT_ID", None) - received_response: dict | None = None - access_token_payload: dict | None = None - sso_assertion: SSOIdentityAssertion | None = None # get url from request if master_key is None: raise ProxyException( @@ -2187,33 +2245,15 @@ async def auth_callback(request: Request, state: str | None = None): redirect_url = SSOAuthenticationHandler.get_redirect_url_for_sso(request=request, sso_callback_route="sso/callback") verbose_proxy_logger.info("Redirecting to %s", redirect_url) - result = None - if google_client_id is not None: - result = await GoogleSSOHandler.get_google_callback_response( - request=request, - google_client_id=google_client_id, - redirect_url=redirect_url, - ) - elif microsoft_client_id is not None: - result = await MicrosoftSSOHandler.get_microsoft_callback_response( - request=request, - microsoft_client_id=microsoft_client_id, - redirect_url=redirect_url, - ) - - elif generic_client_id is not None: - ( - result, - received_response, - access_token_payload, - sso_assertion, - ) = await get_generic_sso_response( - request=request, - jwt_handler=jwt_handler, - generic_client_id=generic_client_id, - redirect_url=redirect_url, - sso_jwt_handler=sso_jwt_handler, - ) + result, received_response, access_token_payload, sso_assertion = await _sso_provider_callback( + request, + redirect_url, + google_client_id=google_client_id, + microsoft_client_id=microsoft_client_id, + generic_client_id=generic_client_id, + jwt_handler=jwt_handler, + sso_jwt_handler=sso_jwt_handler, + ) if result is None: raise HTTPException( @@ -2915,8 +2955,8 @@ def _is_same_origin_return_path(return_to: str) -> bool: async def sso_return_to_redirect( return_to: str | None, jwt_token: str, - redis_usage_cache, - user_api_key_cache, + redis_usage_cache: "RedisCache | None", + user_api_key_cache: UserApiKeyCache, request: Request, ) -> RedirectResponse | None: """Resolve the post-SSO redirect for a ``return_to``, or None to fall through to the dashboard. @@ -3063,6 +3103,7 @@ class SSOAuthenticationHandler: generic_client_id: str | None = None, state: str | None = None, request: Request | None = None, + retention_enabled: RetentionCheck | None = None, ) -> RedirectResponse | None: """ Step 1. Call Get Login Redirect for the SSO provider. Send the redirect response to `redirect_url` @@ -3074,10 +3115,13 @@ class SSOAuthenticationHandler: generic_client_id (Optional[str], optional): The Generic Client ID. Defaults to None. request: Optional FastAPI request, used to drive the ``Secure`` attribute on the ``litellm_oauth_state`` CSRF cookie. + retention_enabled: Whether an ``oauth2_id_jag`` server exists, deciding if the Google and + Microsoft logins ask the IdP for a refresh token. Defaults to the store's own check. Returns: RedirectResponse: The redirect response from the SSO provider. """ + asks_for_refresh_token: Final = await login_asks_for_refresh_token(retention_enabled) # Google SSO Auth if google_client_id is not None: google_client_secret: Final = os.getenv("GOOGLE_CLIENT_SECRET", None) @@ -3099,7 +3143,10 @@ class SSOAuthenticationHandler: google_client_id, ) with google_sso: - return await google_sso.get_login_redirect(state=state) + return await google_sso.get_login_redirect( + state=state, + params=dict(GOOGLE_REFRESH_TOKEN_PARAMS) if asks_for_refresh_token else None, + ) # Microsoft SSO Auth elif microsoft_client_id is not None: microsoft_client_secret: Final = os.getenv("MICROSOFT_CLIENT_SECRET", None) @@ -3117,6 +3164,7 @@ class SSOAuthenticationHandler: tenant=microsoft_tenant, redirect_uri=redirect_url, allow_insecure_http=True, + scope=[*CustomMicrosoftSSO.scope, MICROSOFT_REFRESH_TOKEN_SCOPE] if asks_for_refresh_token else None, ) with microsoft_sso: return await microsoft_sso.get_login_redirect(state=state) @@ -4399,6 +4447,21 @@ class SSOAuthenticationHandler: return userinfo +class _SSOTokenResponse(BaseModel): + """The id_token and refresh token fastapi-sso keeps on the SSO instance after the code exchange, + read through a validated view because the package ships no type information.""" + + model_config = ConfigDict(frozen=True, from_attributes=True) + + id_token: str | None = None + refresh_token: str | None = None + + @field_validator("id_token", "refresh_token", mode="before") + @classmethod + def _strings_only(cls, value: object) -> str | None: + return value if isinstance(value, str) else None + + class MicrosoftSSOHandler: """ Handles Microsoft SSO callback response and returns a CustomOpenID object @@ -4433,9 +4496,10 @@ class MicrosoftSSOHandler: microsoft_client_id: str, redirect_url: str, return_raw_sso_response: bool = False, - ) -> CustomOpenID | OpenID | dict: + ) -> tuple[CustomOpenID | OpenID | dict[str, object], SSOIdentityAssertion | None]: """ - Get the Microsoft SSO callback response + Get the Microsoft SSO callback response, with the IdP identity assertion the token + response carried (the id_token and refresh token an ``oauth2_id_jag`` server spends) Args: return_raw_sso_response: If True, return the raw SSO response @@ -4470,6 +4534,8 @@ class MicrosoftSSOHandler: ) or {} ) + microsoft_tokens: Final = _SSOTokenResponse.model_validate(microsoft_sso) + sso_assertion: Final = assertion_from_sso_login(microsoft_tokens.id_token, microsoft_tokens.refresh_token) if not return_raw_sso_response and "error" in original_msft_result and original_msft_result["error"]: verbose_proxy_logger.warning("Microsoft SSO login failed: Graph /me returned an error instead of a profile") raise HTTPException( @@ -4494,7 +4560,7 @@ class MicrosoftSSOHandler: if return_raw_sso_response: original_msft_result[MicrosoftSSOHandler.GRAPH_API_RESPONSE_KEY] = user_team_ids original_msft_result["app_roles"] = app_roles - return original_msft_result or {} + return original_msft_result or {}, sso_assertion from litellm.proxy.management_endpoints.sso.agent_subject_enrollment import microsoft_interactive_subject @@ -4517,7 +4583,7 @@ class MicrosoftSSOHandler: team_ids=user_team_ids, user_role=user_role, ) - return result + return result, sso_assertion @staticmethod def openid_from_response( @@ -4793,9 +4859,10 @@ class GoogleSSOHandler: google_client_id: str, redirect_url: str, return_raw_sso_response: bool = False, - ) -> OpenID | dict: + ) -> tuple[OpenID | dict[str, object], SSOIdentityAssertion | None]: """ - Get the Google SSO callback response + Get the Google SSO callback response, with the IdP identity assertion the token + response carried (the id_token and refresh token an ``oauth2_id_jag`` server spends) Args: return_raw_sso_response: If True, return the raw SSO response @@ -4816,17 +4883,13 @@ class GoogleSSOHandler: ) # if user is trying to get the raw sso response for debugging, return the raw sso response - if return_raw_sso_response: - return ( - await google_sso.verify_and_process( - request=request, - convert_response=False, - ) - or {} - ) - - result: Final = await google_sso.verify_and_process(request) - return result or {} + result: Final = ( + await google_sso.verify_and_process(request=request, convert_response=False) + if return_raw_sso_response + else await google_sso.verify_and_process(request) + ) + google_tokens: Final = _SSOTokenResponse.model_validate(google_sso) + return result or {}, assertion_from_sso_login(google_tokens.id_token, google_tokens.refresh_token) def _raise_if_sso_debug_disabled() -> None: @@ -4928,37 +4991,16 @@ async def debug_sso_callback(request: Request): else: redirect_url += "/sso/debug/callback" - result = None - received_response: dict | None = None - access_token_payload: dict | None = None - if google_client_id is not None: - result = await GoogleSSOHandler.get_google_callback_response( - request=request, - google_client_id=google_client_id, - redirect_url=redirect_url, - return_raw_sso_response=True, - ) - elif microsoft_client_id is not None: - result = await MicrosoftSSOHandler.get_microsoft_callback_response( - request=request, - microsoft_client_id=microsoft_client_id, - redirect_url=redirect_url, - return_raw_sso_response=True, - ) - - elif generic_client_id is not None: - ( - result, - received_response, - access_token_payload, - _sso_assertion, - ) = await get_generic_sso_response( - request=request, - jwt_handler=jwt_handler, - generic_client_id=generic_client_id, - redirect_url=redirect_url, - sso_jwt_handler=sso_jwt_handler, - ) + result, received_response, access_token_payload, _sso_assertion = await _sso_provider_callback( + request, + redirect_url, + google_client_id=google_client_id, + microsoft_client_id=microsoft_client_id, + generic_client_id=generic_client_id, + jwt_handler=jwt_handler, + sso_jwt_handler=sso_jwt_handler, + return_raw_sso_response=True, + ) # If result is None, return a basic error message if result is None: diff --git a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_refresher.py b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_refresher.py index d80913f8d33..4197303fdd0 100644 --- a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_refresher.py +++ b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_refresher.py @@ -694,6 +694,9 @@ async def test_the_in_process_coordinator_is_retried_until_redis_appears(): {"GENERIC_TOKEN_ENDPOINT": TOKEN_ENDPOINT, "GENERIC_CLIENT_SECRET": "s"}, {"GENERIC_TOKEN_ENDPOINT": TOKEN_ENDPOINT, "GENERIC_CLIENT_ID": "litellm"}, {"GENERIC_TOKEN_ENDPOINT": "", "GENERIC_CLIENT_ID": "litellm", "GENERIC_CLIENT_SECRET": "s"}, + {"GOOGLE_CLIENT_ID": "google-cid"}, + {"MICROSOFT_CLIENT_ID": "ms-cid", "MICROSOFT_CLIENT_SECRET": "s"}, + {"MICROSOFT_CLIENT_ID": "ms-cid", "MICROSOFT_TENANT": "tenant"}, ], ) def test_a_partial_sso_client_is_no_client(env): @@ -717,6 +720,64 @@ def test_the_configured_sso_client_is_the_one_the_login_used(): assert config.client_secret.get_secret_value() == "s3cret" +def test_a_google_login_is_renewed_at_googles_token_endpoint_as_the_google_client(): + """fastapi-sso's Google login redeems the code at Google's token endpoint with HTTP Basic, so + that is where and how the refresh token it handed back is redeemed.""" + config = sso_client_config({"GOOGLE_CLIENT_ID": "google-cid", "GOOGLE_CLIENT_SECRET": "google-secret"}) + + assert config is not None + assert config.token_endpoint == "https://oauth2.googleapis.com/token" + assert config.client_id == "google-cid" + assert config.client_secret.get_secret_value() == "google-secret" + assert config.auth_method == "client_secret_basic" + + +def test_a_microsoft_login_is_renewed_at_the_tenant_token_endpoint_as_the_microsoft_client(): + config = sso_client_config( + {"MICROSOFT_CLIENT_ID": "ms-cid", "MICROSOFT_CLIENT_SECRET": "ms-secret", "MICROSOFT_TENANT": "tenant-123"} + ) + + assert config is not None + assert config.token_endpoint == "https://login.microsoftonline.com/tenant-123/oauth2/v2.0/token" + assert config.client_id == "ms-cid" + assert config.client_secret.get_secret_value() == "ms-secret" + assert config.auth_method == "client_secret_basic" + + +def test_a_custom_microsoft_token_endpoint_wins_over_the_tenant_template(): + """``MICROSOFT_TOKEN_ENDPOINT`` is what the Microsoft login itself honours for sovereign + clouds, so the renewal has to go to the same place.""" + config = sso_client_config( + { + "MICROSOFT_CLIENT_ID": "ms-cid", + "MICROSOFT_CLIENT_SECRET": "ms-secret", + "MICROSOFT_TENANT": "tenant-123", + "MICROSOFT_TOKEN_ENDPOINT": "https://login.microsoftonline.us/tenant-123/oauth2/v2.0/token", + } + ) + + assert config is not None + assert config.token_endpoint == "https://login.microsoftonline.us/tenant-123/oauth2/v2.0/token" + + +def test_the_client_follows_the_login_callbacks_provider_order(): + """A deployment carrying a Google and a generic client signs users in through Google, so the + stored refresh token belongs to the Google client and only that client can redeem it.""" + config = sso_client_config( + { + "GOOGLE_CLIENT_ID": "google-cid", + "GOOGLE_CLIENT_SECRET": "google-secret", + "GENERIC_TOKEN_ENDPOINT": TOKEN_ENDPOINT, + "GENERIC_CLIENT_ID": "litellm", + "GENERIC_CLIENT_SECRET": "s3cret", + } + ) + + assert config is not None + assert config.client_id == "google-cid" + assert config.token_endpoint == "https://oauth2.googleapis.com/token" + + def test_the_live_store_renews_over_the_database_reader(): """The composition root has to produce a renewing store, or none of this runs in production.""" assert isinstance(default_sso_assertion_store(), RefreshingSSOAssertionStore) diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 536c1eb0b44..4961ca0cf42 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -1641,9 +1641,9 @@ class TestMCPServerManager: assert server.oauth2_flow is None @pytest.mark.asyncio - async def test_load_servers_from_config_warns_for_id_jag_with_google_sso(self, monkeypatch, caplog): + async def test_load_servers_from_config_warns_for_id_jag_with_saml_sso(self, monkeypatch, caplog): self._clear_sso_env(monkeypatch) - monkeypatch.setenv("GOOGLE_CLIENT_ID", "google-cid") + monkeypatch.setenv("SAML_IDP_METADATA_URL", "https://idp.example.com/metadata") manager = MCPServerManager() with ( patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=None)), @@ -1655,7 +1655,26 @@ class TestMCPServerManager: warnings = [message for message in caplog.messages if "oauth2_id_jag" in message] assert len(warnings) == 1 assert "idjag_server" in warnings[0] - assert "GENERIC_CLIENT_ID" in warnings[0] + assert "saml" in warnings[0] + assert "MICROSOFT_CLIENT_ID" in warnings[0] + + @pytest.mark.asyncio + @pytest.mark.parametrize("client_id_env", ["GOOGLE_CLIENT_ID", "MICROSOFT_CLIENT_ID", "GENERIC_CLIENT_ID"]) + async def test_load_servers_from_config_is_silent_for_id_jag_under_oidc_logins( + self, monkeypatch, caplog, client_id_env + ): + """Every OIDC login captures the IdP id_token an id_jag server spends, so none of them is a gap.""" + self._clear_sso_env(monkeypatch) + monkeypatch.setenv(client_id_env, "cid") + manager = MCPServerManager() + with ( + patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=None)), + patch.object(manager, "hydrate_config_servers_dcr_clients", new=AsyncMock()), + caplog.at_level(logging.WARNING, logger="LiteLLM"), + ): + await manager.load_servers_from_config(self._id_jag_config()) + + assert not any("oauth2_id_jag" in message for message in caplog.messages) @pytest.mark.asyncio async def test_load_servers_from_config_does_not_warn_for_id_jag_without_sso(self, monkeypatch, caplog): diff --git a/tests/unit/proxy/management_endpoints/test_id_jag_assertion_capture.py b/tests/unit/proxy/management_endpoints/test_id_jag_assertion_capture.py index ff4fbbfb695..2218473de6b 100644 --- a/tests/unit/proxy/management_endpoints/test_id_jag_assertion_capture.py +++ b/tests/unit/proxy/management_endpoints/test_id_jag_assertion_capture.py @@ -49,8 +49,8 @@ class TestActiveSSOProviderMirrorsTheCallback: assert active_sso_provider() is ActiveSSOProvider.none def test_google_outranks_a_co_configured_generic_client(self, monkeypatch): - """The callback tests GOOGLE_CLIENT_ID first, so the generic arm never runs here and - no assertion is captured; reporting generic would clear a gap that is still open.""" + """The callback tests GOOGLE_CLIENT_ID first, so the generic arm never runs here and the + assertion the refresher later renews was minted by the Google client, not the generic one.""" monkeypatch.setenv("GOOGLE_CLIENT_ID", "google-cid") monkeypatch.setenv("GENERIC_CLIENT_ID", "generic-cid") assert active_sso_provider() is ActiveSSOProvider.google @@ -67,23 +67,20 @@ class TestActiveSSOProviderMirrorsTheCallback: class TestIdJagAssertionCaptureGap: - def test_generic_oidc_has_no_gap(self, monkeypatch): - monkeypatch.setenv("GENERIC_CLIENT_ID", "generic-cid") + @pytest.mark.parametrize("env_var", ["GOOGLE_CLIENT_ID", "MICROSOFT_CLIENT_ID", "GENERIC_CLIENT_ID"]) + def test_every_oidc_login_captures_so_has_no_gap(self, monkeypatch, env_var): + """Google, Microsoft and generic OIDC all read the id_token off the token response, so a + deployment on any of them must not be told to switch providers.""" + monkeypatch.setenv(env_var, "configured") assert id_jag_assertion_capture_gap() is None - @pytest.mark.parametrize( - "env_var, provider_label", - [ - ("GOOGLE_CLIENT_ID", "google"), - ("MICROSOFT_CLIENT_ID", "microsoft"), - ("SAML_IDP_METADATA_URL", "saml"), - ], - ) - def test_non_capturing_provider_is_named_with_the_remedy(self, monkeypatch, env_var, provider_label): - monkeypatch.setenv(env_var, "configured") + def test_saml_is_named_with_the_oidc_remedy(self, monkeypatch): + monkeypatch.setenv("SAML_IDP_METADATA_URL", "https://idp.example.com/metadata") gap = id_jag_assertion_capture_gap() assert gap is not None - assert provider_label in gap + assert "saml" in gap + assert "GOOGLE_CLIENT_ID" in gap + assert "MICROSOFT_CLIENT_ID" in gap assert "GENERIC_CLIENT_ID" in gap def test_no_sso_configured_reports_a_gap(self): @@ -91,27 +88,26 @@ class TestIdJagAssertionCaptureGap: assert gap is not None assert "no SSO provider is configured" in gap - def test_google_beside_generic_still_reports_a_gap(self, monkeypatch): - """The precedence trap in operator terms: adding a generic client id without removing - GOOGLE_CLIENT_ID does not fix the deployment, so the gap must not clear.""" - monkeypatch.setenv("GOOGLE_CLIENT_ID", "google-cid") - monkeypatch.setenv("GENERIC_CLIENT_ID", "generic-cid") - gap = id_jag_assertion_capture_gap() - assert gap is not None - assert "google" in gap + def test_saml_beside_an_oidc_client_has_no_gap(self, monkeypatch): + """The callback dispatches to the OIDC client before SAML, so that is the login that runs + and it captures; reporting SAML's gap would tell the operator to fix a path nobody takes.""" + monkeypatch.setenv("MICROSOFT_CLIENT_ID", "ms-cid") + monkeypatch.setenv("SAML_IDP_METADATA_URL", "https://idp.example.com/metadata") + assert id_jag_assertion_capture_gap() is None class TestIdJagAssertionCaptureGapAtStartup: def test_no_provider_at_startup_is_not_yet_a_gap(self): assert id_jag_assertion_capture_gap_at_startup() is None - def test_google_provider_at_startup_reports_the_capture_gap(self, monkeypatch): - monkeypatch.setenv("GOOGLE_CLIENT_ID", "google-cid") + def test_saml_at_startup_reports_the_capture_gap(self, monkeypatch): + monkeypatch.setenv("SAML_IDP_METADATA_URL", "https://idp.example.com/metadata") startup_gap = id_jag_assertion_capture_gap_at_startup() callback_gap = id_jag_assertion_capture_gap() assert startup_gap is not None assert startup_gap == callback_gap - def test_generic_provider_at_startup_has_no_gap(self, monkeypatch): - monkeypatch.setenv("GENERIC_CLIENT_ID", "generic-cid") + @pytest.mark.parametrize("env_var", ["GOOGLE_CLIENT_ID", "MICROSOFT_CLIENT_ID", "GENERIC_CLIENT_ID"]) + def test_an_oidc_provider_at_startup_has_no_gap(self, monkeypatch, env_var): + monkeypatch.setenv(env_var, "configured") assert id_jag_assertion_capture_gap_at_startup() is None diff --git a/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py index 47aea89843f..cd5d5e475dd 100644 --- a/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -4458,9 +4458,9 @@ class TestAddMCPServerAtomicity: class TestIdJagRegistrationWarnsAboutTheSSOGap: """An `oauth2_id_jag` server only ever works when the login path captures an IdP identity - assertion, and only the generic OIDC arm does. Registering one under Google or Microsoft - succeeds and then fails for every user on every call, so the mismatch has to be said at - registration time, while the admin is still looking at the configuration.""" + assertion, and only the OIDC logins (Google, Microsoft, generic) do. Registering one under SAML + or with no SSO succeeds and then fails for every user on every call, so the mismatch has to be + said at registration time, while the admin is still looking at the configuration.""" @staticmethod def _clear_sso_env(monkeypatch): @@ -4530,8 +4530,6 @@ class TestIdJagRegistrationWarnsAboutTheSSOGap: @pytest.mark.parametrize( "provider_env, expected_fragment", [ - ({"GOOGLE_CLIENT_ID": "cid"}, "google"), - ({"MICROSOFT_CLIENT_ID": "cid"}, "microsoft"), ({"SAML_IDP_METADATA_URL": "https://idp.example.com/metadata"}, "saml"), ({}, "no SSO provider is configured"), ], @@ -4546,8 +4544,9 @@ class TestIdJagRegistrationWarnsAboutTheSSOGap: assert "ema-1" in str(warnings[0]) @pytest.mark.asyncio - async def test_create_is_silent_under_generic_oidc(self, monkeypatch, caplog): - await self._run_create(monkeypatch, {"GENERIC_CLIENT_ID": "cid"}, MCPAuth.oauth2_id_jag, caplog) + @pytest.mark.parametrize("client_id_env", ["GOOGLE_CLIENT_ID", "MICROSOFT_CLIENT_ID", "GENERIC_CLIENT_ID"]) + async def test_create_is_silent_under_oidc_logins(self, monkeypatch, caplog, client_id_env): + await self._run_create(monkeypatch, {client_id_env: "cid"}, MCPAuth.oauth2_id_jag, caplog) assert self._id_jag_warnings(caplog) == [] @pytest.mark.asyncio @@ -4565,7 +4564,7 @@ class TestIdJagRegistrationWarnsAboutTheSSOGap: ) self._clear_sso_env(monkeypatch) - monkeypatch.setenv("GOOGLE_CLIENT_ID", "cid") + monkeypatch.setenv("SAML_IDP_METADATA_URL", "https://idp.example.com/metadata") mock_manager = MagicMock() mock_manager.update_server = AsyncMock() @@ -4602,7 +4601,7 @@ class TestIdJagRegistrationWarnsAboutTheSSOGap: warnings = self._id_jag_warnings(caplog) assert len(warnings) == 1 - assert "google" in str(warnings[0]) + assert "saml" in str(warnings[0]) class TestHealthCheckServers: diff --git a/tests/unit/proxy/management_endpoints/test_ui_sso.py b/tests/unit/proxy/management_endpoints/test_ui_sso.py index f69c069cfea..91a8d4f7101 100644 --- a/tests/unit/proxy/management_endpoints/test_ui_sso.py +++ b/tests/unit/proxy/management_endpoints/test_ui_sso.py @@ -243,17 +243,24 @@ def test_get_microsoft_callback_response(stubbed_graph_api): "surname": "User", } + id_token = _ema_id_token(sub="msft123") + + async def fake_verify(self, *args, **kwargs): + self._id_token = id_token + self._refresh_token = "msft-refresh-token" + return mock_response + with patch.dict( os.environ, {"MICROSOFT_CLIENT_SECRET": "mock_secret", "MICROSOFT_TENANT": "mock_tenant"}, ): - mock_verify = AsyncMock(return_value=mock_response) - with patch( + with patch( # test-quality-ok: stubs the upstream Microsoft IdP boundary, where the token response is read "fastapi_sso.sso.microsoft.MicrosoftSSO.verify_and_process", - new=mock_verify, + autospec=True, + side_effect=fake_verify, ): # Act - result = asyncio.run( + result, sso_assertion = asyncio.run( MicrosoftSSOHandler.get_microsoft_callback_response( request=mock_request, microsoft_client_id="mock_client_id", @@ -269,6 +276,11 @@ def test_get_microsoft_callback_response(stubbed_graph_api): assert result.id == "msft123" assert result.first_name == "Microsoft" assert result.last_name == "User" + assert sso_assertion is not None + assert sso_assertion.id_token.get_secret_value() == id_token + assert sso_assertion.refresh_token is not None + assert sso_assertion.refresh_token.get_secret_value() == "msft-refresh-token" + assert sso_assertion.issuer == "https://idp.example.com" GRAPH_ERROR_BODIES: Final = ( @@ -317,7 +329,7 @@ def test_get_microsoft_callback_response_signs_in_a_profile_with_an_empty_error_ "fastapi_sso.sso.microsoft.MicrosoftSSO.verify_and_process", new=AsyncMock(return_value={**profile, "error": empty_error}), ): - result = asyncio.run( + result, _ = asyncio.run( MicrosoftSSOHandler.get_microsoft_callback_response( request=mock_request, microsoft_client_id="mock_client_id", @@ -342,7 +354,7 @@ def test_get_microsoft_callback_response_raw_sso_response_keeps_the_graph_error( "fastapi_sso.sso.microsoft.MicrosoftSSO.verify_and_process", new=AsyncMock(return_value=dict(graph_error)), ): - result = asyncio.run( + result, _ = asyncio.run( MicrosoftSSOHandler.get_microsoft_callback_response( request=mock_request, microsoft_client_id="mock_client_id", @@ -376,7 +388,7 @@ def test_get_microsoft_callback_response_raw_sso_response(stubbed_graph_api): new=mock_verify, ): # Act - result = asyncio.run( + result, _ = asyncio.run( MicrosoftSSOHandler.get_microsoft_callback_response( request=mock_request, microsoft_client_id="mock_client_id", @@ -405,13 +417,19 @@ def test_get_google_callback_response(): "family_name": "User", } + id_token = _ema_id_token(sub="google123") + + async def fake_verify(self, *args, **kwargs): + self._id_token = id_token + self._refresh_token = None + return mock_response + with patch.dict(os.environ, {"GOOGLE_CLIENT_SECRET": "mock_secret"}): - mock_verify = AsyncMock(return_value=mock_response) - with patch( - "fastapi_sso.sso.google.GoogleSSO.verify_and_process", new=mock_verify + with patch( # test-quality-ok: stubs the upstream Google IdP boundary, where the token response is read + "fastapi_sso.sso.google.GoogleSSO.verify_and_process", autospec=True, side_effect=fake_verify ): # Act - result = asyncio.run( + result, sso_assertion = asyncio.run( GoogleSSOHandler.get_google_callback_response( request=mock_request, google_client_id="mock_client_id", @@ -426,6 +444,9 @@ def test_get_google_callback_response(): assert result.get("sub") == "google123" assert result.get("given_name") == "Google" assert result.get("family_name") == "User" + assert sso_assertion is not None + assert sso_assertion.id_token.get_secret_value() == id_token + assert sso_assertion.refresh_token is None @pytest.mark.asyncio @@ -4018,7 +4039,13 @@ class TestCLIKeyRegenerationFlow: cli_state = f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:cli-new-session-key-456" # Mock the CLI callback and required proxy server components + from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( + assertion_from_sso_login, + ) + mock_result = {"user_id": "test-user", "email": "test@example.com"} + google_assertion = assertion_from_sso_login(_ema_id_token(sub="test-user"), None) + assert google_assertion is not None with ( patch( @@ -4032,7 +4059,7 @@ class TestCLIKeyRegenerationFlow: patch.dict(os.environ, {"GOOGLE_CLIENT_ID": "test-google-id"}, clear=True), patch( "litellm.proxy.management_endpoints.ui_sso.GoogleSSOHandler.get_google_callback_response", - return_value=mock_result, + return_value=(mock_result, google_assertion), ), ): mock_cli_callback.return_value = MagicMock() @@ -4046,7 +4073,7 @@ class TestCLIKeyRegenerationFlow: prefill_user_code=None, result=mock_result, received_response=None, - sso_assertion=None, + sso_assertion=google_assertion, ) @pytest.mark.asyncio @@ -4074,7 +4101,7 @@ class TestCLIKeyRegenerationFlow: patch.dict(os.environ, {"GOOGLE_CLIENT_ID": "test-google-id"}, clear=True), patch( "litellm.proxy.management_endpoints.ui_sso.GoogleSSOHandler.get_google_callback_response", - return_value=mock_result, + return_value=(mock_result, None), ), ): mock_cli_callback.return_value = MagicMock() @@ -8836,7 +8863,7 @@ async def test_debug_sso_callback_handles_missing_raw_response(): ) async def fake_microsoft_callback(**kwargs): - return parsed_openid + return parsed_openid, None with ( patch.dict( @@ -8886,7 +8913,7 @@ async def _render_debug_page(provider_env, id_jag_registered, force_inert=False) return parsed, {"sub": "user_123"}, {"scope": "openid"}, None async def fake_google(**kwargs): - return parsed + return parsed, None stack = [ patch.dict(os.environ, {**provider_env, "ENABLE_SSO_DEBUG": "true"}, clear=False), @@ -8925,14 +8952,14 @@ async def _render_debug_page(provider_env, id_jag_registered, force_inert=False) @pytest.mark.asyncio -async def test_debug_page_logs_the_capture_gap_but_never_renders_it(caplog): +async def test_debug_page_logs_no_capture_gap_for_a_google_login(caplog): + """Google reads the id_token off its token response like generic OIDC does, so the debug page + under GOOGLE_CLIENT_ID with an oauth2_id_jag server registered has no gap to log and nothing + to render about it.""" with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): body = await _render_debug_page({"GOOGLE_CLIENT_ID": _GOOGLE_DEBUG_CLIENT_ID}, id_jag_registered=True) - warnings = _id_jag_gap_warnings(caplog) - assert len(warnings) == 1 - assert "google" in warnings[0] - assert "GENERIC_CLIENT_ID" in warnings[0] + assert _id_jag_gap_warnings(caplog) == [] assert "id_jag" not in body assert "GENERIC_CLIENT_ID" not in body @@ -8977,7 +9004,7 @@ async def test_debug_page_survives_a_store_outage(monkeypatch, caplog): not the page.""" from litellm.proxy.management_endpoints.ui_sso import warn_if_id_jag_capture_gap - monkeypatch.setenv("GOOGLE_CLIENT_ID", _GOOGLE_DEBUG_CLIENT_ID) + monkeypatch.setenv("SAML_IDP_METADATA_URL", "https://idp.example.com/metadata") retention_check = AsyncMock(side_effect=Exception("db down")) with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): assert await warn_if_id_jag_capture_gap(retention_enabled=retention_check) is None @@ -9635,6 +9662,158 @@ async def test_redirect_from_openid_persists_assertion_under_canonical_user_id() assert response is not None +@pytest.mark.asyncio +@pytest.mark.parametrize( + "provider_env, sso_class, payload, needs_graph", + [ + ( + {"MICROSOFT_CLIENT_ID": "ms-cid", "MICROSOFT_CLIENT_SECRET": "ms-secret", "MICROSOFT_TENANT": "tenant"}, + "fastapi_sso.sso.microsoft.MicrosoftSSO", + {"mail": "u@example.com", "displayName": "U Ser", "id": "idp-subject"}, + True, + ), + ( + {"GOOGLE_CLIENT_ID": "google-cid", "GOOGLE_CLIENT_SECRET": "google-secret"}, + "fastapi_sso.sso.google.GoogleSSO", + {"email": "u@example.com", "name": "U Ser", "sub": "idp-subject"}, + False, + ), + ], +) +async def test_ui_login_through_microsoft_or_google_hands_its_assertion_to_the_redirect( + request, provider_env, sso_class, payload, needs_graph +): + """The UI funnel used to capture an assertion only on the generic branch, so an ``oauth2_id_jag`` + server had nothing to spend for a deployment signing in through Entra or Google. The real + Microsoft and Google handlers run here, only the IdP token exchange is stubbed, so the id_token + and refresh token the IdP returned must reach the redirect that persists them.""" + from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler, auth_callback + + if needs_graph: + request.getfixturevalue("stubbed_graph_api") + id_token = _ema_id_token(sub="idp-subject") + + async def fake_verify(self, *args, **kwargs): + self._id_token = id_token + self._refresh_token = "idp-refresh" + return payload + + mock_request = MagicMock(spec=Request) + mock_request.query_params = {"code": "auth-code"} + mock_request.cookies = {} + mock_request.scope = {} + mock_request.base_url = "http://localhost:4000/" + redirect_mock = AsyncMock(return_value=MagicMock()) + + with ( + patch.dict(os.environ, provider_env, clear=True), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.master_key", "sk-master"), + patch("litellm.proxy.proxy_server.general_settings", {}), + patch("litellm.proxy.proxy_server.jwt_handler", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + patch( # test-quality-ok: stubs the upstream IdP token exchange, where the token response is read + f"{sso_class}.verify_and_process", autospec=True, side_effect=fake_verify + ), + patch.object( # test-quality-ok: the DB and key-mint boundary after the login, asserted on its inputs + SSOAuthenticationHandler, "get_redirect_response_from_openid", redirect_mock + ), + ): + await auth_callback(request=mock_request, state=None) + + redirect_mock.assert_awaited_once() + assert redirect_mock.await_args.kwargs["result"] is not None + sso_assertion = redirect_mock.await_args.kwargs["sso_assertion"] + assert sso_assertion is not None + assert sso_assertion.id_token.get_secret_value() == id_token + assert sso_assertion.refresh_token is not None + assert sso_assertion.refresh_token.get_secret_value() == "idp-refresh" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "retention_enabled, expected_params", + [ + (AsyncMock(return_value=True), {"access_type": "offline", "prompt": "consent"}), + (AsyncMock(return_value=False), None), + (AsyncMock(side_effect=RuntimeError("store down")), None), + ], +) +async def test_google_login_asks_for_a_refresh_token_only_while_an_id_jag_server_exists( + retention_enabled, expected_params +): + """Google issues a refresh token only when the authorization request carries ``access_type=offline``, + and the refresher cannot renew a Google assertion without one. The login asks only while an + ``oauth2_id_jag`` server exists, and a store outage must not fail the login or add the consent.""" + from fastapi.responses import RedirectResponse + + from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler + + seen: dict[str, object] = {} + + async def fake_redirect(self, *, redirect_uri=None, params=None, state=None): + seen["params"] = params + return RedirectResponse(url="https://accounts.google.com/o/oauth2/v2/auth") + + with ( + patch.dict(os.environ, {"GOOGLE_CLIENT_SECRET": "google-secret"}), + patch( # test-quality-ok: stubs the upstream Google authorization boundary, where the request leaves + "fastapi_sso.sso.google.GoogleSSO.get_login_redirect", autospec=True, side_effect=fake_redirect + ), + ): + response = await SSOAuthenticationHandler.get_sso_login_redirect( + redirect_url="http://localhost:4000/sso/callback", + google_client_id="google-cid", + retention_enabled=retention_enabled, + ) + + assert response is not None + assert seen["params"] == expected_params + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "retention_enabled, asks_offline_access", + [ + (AsyncMock(return_value=True), True), + (AsyncMock(return_value=False), False), + (AsyncMock(side_effect=RuntimeError("store down")), False), + ], +) +async def test_microsoft_login_asks_for_offline_access_only_while_an_id_jag_server_exists( + retention_enabled, asks_offline_access +): + """Entra issues a refresh token only when ``offline_access`` is in the requested scope, and the + refresher cannot renew an Entra assertion without one. The default sign-in scopes stay in place so + the login keeps resolving the same profile.""" + from fastapi.responses import RedirectResponse + from fastapi_sso.sso.microsoft import MicrosoftSSO + + from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler + + seen: dict[str, object] = {} + + async def fake_redirect(self, *, redirect_uri=None, params=None, state=None): + seen["scope"] = list(self._scope) + return RedirectResponse(url="https://login.microsoftonline.com/tenant/oauth2/v2.0/authorize") + + with ( + patch.dict(os.environ, {"MICROSOFT_CLIENT_SECRET": "ms-secret", "MICROSOFT_TENANT": "tenant"}), + patch( # test-quality-ok: stubs the upstream Entra authorization boundary, where the request leaves + "fastapi_sso.sso.microsoft.MicrosoftSSO.get_login_redirect", autospec=True, side_effect=fake_redirect + ), + ): + response = await SSOAuthenticationHandler.get_sso_login_redirect( + redirect_url="http://localhost:4000/sso/callback", + microsoft_client_id="ms-cid", + retention_enabled=retention_enabled, + ) + + assert response is not None + assert all(default in seen["scope"] for default in MicrosoftSSO.scope), seen["scope"] + assert ("offline_access" in seen["scope"]) is asks_offline_access, seen["scope"] + + @pytest.mark.asyncio async def test_cli_completion_persists_assertion_under_db_user_id(): """The CLI funnel persists the captured assertion under the DB-resolved user_id.""" @@ -9710,8 +9889,7 @@ def _id_jag_gap_warnings(caplog) -> list[str]: @pytest.mark.parametrize( "provider_env, expected_fragment", [ - ({"GOOGLE_CLIENT_ID": "cid"}, "google"), - ({"MICROSOFT_CLIENT_ID": "cid", "MICROSOFT_TENANT": "t"}, "microsoft"), + ({"SAML_IDP_METADATA_URL": "https://idp.example.com/metadata"}, "saml"), ({}, "no SSO provider is configured"), ], ) @@ -9738,16 +9916,17 @@ async def test_uncaptured_assertion_warns_when_an_id_jag_server_is_registered( @pytest.mark.asyncio -async def test_generic_provider_that_returned_no_id_token_still_warns(monkeypatch, caplog): - """Generic OIDC has a capture path, so there is no configuration gap to report; the login - still handed the id_jag arm nothing, and that must not pass silently.""" +@pytest.mark.parametrize("provider_env_var", ["GOOGLE_CLIENT_ID", "MICROSOFT_CLIENT_ID", "GENERIC_CLIENT_ID"]) +async def test_a_capturing_provider_that_returned_no_id_token_still_warns(monkeypatch, caplog, provider_env_var): + """Google, Microsoft and generic OIDC all have a capture path, so there is no configuration gap + to report; the login still handed the id_jag arm nothing, and that must not pass silently.""" from litellm.proxy.management_endpoints.ui_sso import ( warn_if_id_jag_assertion_uncaptured, ) - for name in ("GOOGLE_CLIENT_ID", "MICROSOFT_CLIENT_ID", "SAML_IDP_METADATA_URL"): + for name in ("GOOGLE_CLIENT_ID", "MICROSOFT_CLIENT_ID", "GENERIC_CLIENT_ID", "SAML_IDP_METADATA_URL"): monkeypatch.delenv(name, raising=False) - monkeypatch.setenv("GENERIC_CLIENT_ID", "cid") + monkeypatch.setenv(provider_env_var, "cid") with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): await warn_if_id_jag_assertion_uncaptured(None, retention_enabled=AsyncMock(return_value=True)) @@ -9875,7 +10054,7 @@ async def test_browser_funnel_reports_an_uncaptured_assertion(monkeypatch, caplo warnings = _id_jag_gap_warnings(caplog) assert len(warnings) == 1 - assert "google" in str(warnings[0]) + assert "no usable id_token" in str(warnings[0]) @pytest.mark.asyncio @@ -9939,7 +10118,7 @@ async def test_cli_funnel_reports_an_uncaptured_assertion(monkeypatch, caplog): warnings = _id_jag_gap_warnings(caplog) assert len(warnings) == 1 - assert "microsoft" in str(warnings[0]) + assert "no usable id_token" in str(warnings[0]) def _cli_callback_kwargs(flow): From f02449ad3801c1a638fb2cd03987179fcfbe0712 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Fri, 16 Oct 2026 12:16:56 -0700 Subject: [PATCH 14/58] feat(sso): retain the SAML assertion at login as the saml2 subject token for oauth2_id_jag (#46745) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/mcp_server_manager.py | 3 + .../outbound_credentials/__init__.py | 2 + .../outbound_credentials/adapter.py | 17 +- .../login_assertion_exchange.py | 216 +++++++++++ .../outbound_credentials/resolver.py | 79 +++- .../sso_assertion_refresher.py | 19 +- .../sso_assertion_store.py | 142 ++++++- .../mcp_server/outbound_credentials/types.py | 9 + litellm/proxy/_lazy_openapi_snapshot.json | 44 +++ .../sso/id_jag_assertion_capture.py | 22 +- .../management_endpoints/sso/saml_sso.py | 58 ++- litellm/proxy/management_endpoints/ui_sso.py | 14 +- litellm/types/mcp.py | 9 + .../types/mcp_server/mcp_server_manager.py | 1 + .../outbound_credentials/test_adapter.py | 23 ++ .../test_login_assertion_exchange.py | 327 ++++++++++++++++ .../outbound_credentials/test_resolver.py | 354 +++++++++++++++--- .../test_sso_assertion_refresher.py | 89 +++++ .../test_sso_assertion_store.py | 152 +++++++- .../mcp_server/test_mcp_server_manager.py | 9 +- .../test_id_jag_assertion_capture.py | 18 +- .../test_mcp_management_endpoints.py | 6 +- .../management_endpoints/test_saml_sso.py | 221 +++++++---- .../proxy/management_endpoints/test_ui_sso.py | 84 ++++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 + 25 files changed, 1720 insertions(+), 200 deletions(-) create mode 100644 litellm/proxy/_experimental/mcp_server/outbound_credentials/login_assertion_exchange.py create mode 100644 tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_login_assertion_exchange.py diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index fe34e18aecd..b724a61cd94 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -455,6 +455,7 @@ class MCPServerConfig(TypedDict, total=False): upstream_token_header: ReadOnly[str] id_jag_resource_token_endpoint: str id_jag_resource: str + id_jag_subject_source: ReadOnly[str] id_jag_resource_client_id: ReadOnly[str] id_jag_resource_client_secret: ReadOnly[str] client_private_key: str @@ -2789,6 +2790,7 @@ class MCPServerManager: # ID-JAG fields id_jag_resource_token_endpoint=server_config.get("id_jag_resource_token_endpoint", None), id_jag_resource=server_config.get("id_jag_resource", None), + id_jag_subject_source=server_config.get("id_jag_subject_source", None), id_jag_resource_client_id=server_config.get("id_jag_resource_client_id", None), id_jag_resource_client_secret=server_config.get("id_jag_resource_client_secret", None), client_private_key=server_config.get("client_private_key", None), @@ -3380,6 +3382,7 @@ class MCPServerManager: credentials_dict.get("id_jag_resource_token_endpoint") if credentials_dict else None ), id_jag_resource=(credentials_dict.get("id_jag_resource") if credentials_dict else None), + id_jag_subject_source=(credentials_dict.get("id_jag_subject_source") if credentials_dict else None), id_jag_resource_client_id=(credentials_dict.get("id_jag_resource_client_id") if credentials_dict else None), id_jag_resource_client_secret=self._decrypt_credential_field( credentials_dict.get("id_jag_resource_client_secret") if credentials_dict else None, diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/__init__.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/__init__.py index 20fac28a55b..3260603c3ca 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/__init__.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/__init__.py @@ -39,6 +39,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( HeaderCarrier, IdJagConfig, IdJagResourceClient, + IdJagSubjectSource, NoneConfig, PassthroughConfig, PrivateKeyJwtAuth, @@ -71,6 +72,7 @@ __all__ = [ "HeaderCarrier", "IdJagConfig", "IdJagResourceClient", + "IdJagSubjectSource", "NoOpAuth", "NoneConfig", "Ok", diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py index 935d026b22c..f217072bd39 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py @@ -34,6 +34,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( CredError, IdJagConfig, IdJagResourceClient, + IdJagSubjectSource, NoneConfig, PassthroughConfig, PrivateKeyJwtAuth, @@ -236,7 +237,8 @@ def _id_jag_spec(server: MCPServer, resource: str) -> ServerSpec | None: if not org_token_endpoint or not resource_token_endpoint or not client_id or client_auth is None: return None resource_client: Final = _id_jag_resource_client(server, client_id, client_auth) - if resource_client is None: + subject_source: Final = _id_jag_subject_source(server) + if resource_client is None or subject_source is None: return None return ServerSpec( server_id=server.server_id, @@ -249,6 +251,7 @@ def _id_jag_spec(server: MCPServer, resource: str) -> ServerSpec | None: client_auth=client_auth, resource_client=resource_client, subject_token_type=_id_jag_subject_token_type(server), + subject_source=subject_source, audience=server.audience, resource=server.id_jag_resource, scopes=tuple(server.scopes or ()), @@ -282,6 +285,18 @@ def _id_jag_client_auth(server: MCPServer) -> ClientAuth | None: return None +def _id_jag_subject_source(server: MCPServer) -> IdJagSubjectSource | None: + """``None`` for a value that names no source, so the half-configured server defers like a missing + endpoint does instead of silently presenting the wrong token.""" + configured: Final = server.id_jag_subject_source + if not configured: + return IdJagSubjectSource.login_assertion + try: + return IdJagSubjectSource(configured) + except ValueError: + return None + + def _id_jag_subject_token_type(server: MCPServer) -> str: """ID-JAG asserts the user's id_token, so the token-exchange access_token default maps to id_token; an explicitly configured value (e.g. a SAML2 assertion type) is honored verbatim.""" diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/login_assertion_exchange.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/login_assertion_exchange.py new file mode 100644 index 00000000000..224f956bcfd --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/login_assertion_exchange.py @@ -0,0 +1,216 @@ +"""Exchange the identity assertion captured at SSO login for a refresh token, at login. + +An ``oauth2_id_jag`` server configured with ``id_jag_subject_source: refresh_token`` does not present the +login assertion on ID-JAG leg 1. It presents a refresh token that the IdP issued for that assertion through +one RFC 8693 exchange (``requested_token_type`` refresh_token), which is how a SAML requesting app works at +an IdP whose ID-JAG exchange does not take a ``saml2`` subject directly. That exchange has to happen at +login: a SAML assertion is valid for minutes, the first MCP call that needs it can come hours later, and +the refresh token outlives the assertion by as long as the IdP decides. + +The servers consulted are the two configuration authorities, the pod-local config declaration and the DB +rows, like the retention gate, not the in-memory registry, which is a per-process snapshot that can lag a +server added on another pod. The exchange authenticates as the ID-JAG server's own IdP client, since the +refresh token is bound to the client that later spends it. One row holds one refresh token per user, so every server in this mode has +to share one IdP client; a second client is logged and skipped rather than silently handed a token it +cannot redeem. A refused or unreachable IdP never fails the login, the row keeps the plain assertion and +the ID-JAG call reports the missing refresh token as the precondition it is. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import Iterable, Mapping, Sequence +from datetime import datetime, timezone +from typing import Final, Protocol + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok, Result +from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( + SSOIdentityAssertion, + with_login_refresh_token, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.token_endpoint import ExchangedToken +from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( + ClientAuth, + CredError, + IdJagConfig, + IdJagSubjectSource, + ServerSpec, +) +from litellm.proxy._types import LiteLLM_MCPServerTable +from litellm.types.mcp import MCPAuth +from litellm.types.mcp_server.mcp_server_manager import MCPServer + +_TOKEN_EXCHANGE_GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:token-exchange" +REFRESH_TOKEN_TOKEN_TYPE: Final = "urn:ietf:params:oauth:token-type:refresh_token" +_LOGIN_EXCHANGE_DEADLINE_SECONDS: Final = 15.0 + +_REFRESH_TOKEN_SCOPE: Final = "openid offline_access" +_RUNTIME_APPROVAL_STATES: Final = (None, "active", "approved") # the states get_runtime_mcp_server_rows serves + + +class LoginExchangeCatalog(Protocol): + """The slice of the MCP server manager the login-time exchange reads.""" + + @property + def config_mcp_servers(self) -> Mapping[str, MCPServer]: ... + + async def build_mcp_server_from_table( + self, + mcp_server: LiteLLM_MCPServerTable, + *, + credentials_are_encrypted: bool = True, + env_vars_are_encrypted: bool | None = None, + register_oauth_discovery: bool = True, + ) -> MCPServer: ... + + +class TokenExchangeEndpoint(Protocol): + async def fetch( + self, endpoint: str, client_id: str, grant_params: Mapping[str, str], client_auth: ClientAuth + ) -> Result[ExchangedToken, CredError]: ... + + +def _is_runtime_id_jag_row(row: LiteLLM_MCPServerTable) -> bool: + """An ID-JAG row the runtime registry would serve. A pending or rejected submission never gets a + login assertion posted to the token endpoint it names.""" + return row.auth_type == MCPAuth.oauth2_id_jag and row.approval_status in _RUNTIME_APPROVAL_STATES + + +async def login_exchange_servers( + catalog: LoginExchangeCatalog, db_rows: Sequence[LiteLLM_MCPServerTable] +) -> tuple[MCPServer, ...]: + """Every server the exchange has to consider: the config-declared ones plus the approved ID-JAG rows + in the DB, built the same way the catalog builds them so their credentials are decrypted. A DB row + hides a config server with the same ``server_id``, exactly as the live catalog serves them.""" + id_jag_rows: Final = tuple(row for row in db_rows if _is_runtime_id_jag_row(row)) + db_servers: Final = tuple([await _built_or_skipped(catalog, row) for row in id_jag_rows]) + db_by_id: Final = {server.server_id: server for server in db_servers if server is not None} + return tuple({**catalog.config_mcp_servers, **db_by_id}.values()) + + +async def _built_or_skipped(catalog: LoginExchangeCatalog, row: LiteLLM_MCPServerTable) -> MCPServer | None: + """One row that cannot be built (undecryptable credentials, a half-configured submission) is skipped + with a warning. It must not cost every other refresh-token server its login exchange.""" + try: + return await catalog.build_mcp_server_from_table(row, register_oauth_discovery=False) + except Exception as exc: # noqa: BLE001 # a per-row build failure is logged and isolated, not propagated + verbose_proxy_logger.warning( + "ID-JAG: MCP server row server_id=%s could not be built for the login-time refresh token exchange " + "and is skipped: %s", + row.server_id, + exc, + ) + return None + + +def refresh_token_exchange_targets(specs: Iterable[ServerSpec | None]) -> tuple[IdJagConfig, ...]: + """The ID-JAG configs that present a login refresh token, one per IdP client, in registration order.""" + configs: Final = tuple( + spec.config + for spec in specs + if spec is not None + and isinstance(spec.config, IdJagConfig) + and spec.config.subject_source is IdJagSubjectSource.refresh_token + ) + clients: Final = tuple((config.org_token_endpoint, config.client_id) for config in configs) + return tuple(config for index, config in enumerate(configs) if clients.index(clients[index]) == index) + + +async def exchange_login_assertion_for_refresh_token( + assertion: SSOIdentityAssertion, + targets: tuple[IdJagConfig, ...], + token_endpoint: TokenExchangeEndpoint, + now: datetime, + deadline_seconds: float = _LOGIN_EXCHANGE_DEADLINE_SECONDS, +) -> SSOIdentityAssertion: + """The assertion as captured when no server wants a refresh token, otherwise the assertion plus the + refresh token the first target's IdP client exchanged it for. An OIDC login's own SSO refresh token + does not count, it is bound to the SSO client and the ID-JAG client cannot redeem it.""" + if not targets or assertion.exchanged_refresh_token is not None: + return assertion + config: Final = targets[0] + if len(targets) > 1: + verbose_proxy_logger.warning( + "ID-JAG: %d IdP clients are configured with id_jag_subject_source refresh_token; the login refresh " + "token is minted for client_id=%s at %s and servers using another client cannot redeem it.", + len(targets), + config.client_id, + config.org_token_endpoint, + ) + params: Final = { + "grant_type": _TOKEN_EXCHANGE_GRANT_TYPE, + "requested_token_type": REFRESH_TOKEN_TOKEN_TYPE, + "subject_token": assertion.id_token.get_secret_value(), + "subject_token_type": assertion.subject_token_type or config.subject_token_type, + "scope": _REFRESH_TOKEN_SCOPE, + } + try: + outcome: Final = await asyncio.wait_for( + token_endpoint.fetch(config.org_token_endpoint, config.client_id, params, config.client_auth), + timeout=deadline_seconds, + ) + except (TimeoutError, asyncio.TimeoutError): + verbose_proxy_logger.warning( + "ID-JAG: the IdP at %s did not answer the login-time refresh token exchange within %.0fs, so the login " + "finishes with the captured assertion only. Servers with id_jag_subject_source: refresh_token will ask " + "this user to sign in again.", + config.org_token_endpoint, + deadline_seconds, + ) + return assertion + match outcome: + case Error(err): + verbose_proxy_logger.warning( + "ID-JAG: the IdP at %s refused to exchange the login assertion for a refresh token (%s); the " + "login goes on with the assertion alone and servers presenting a refresh token will ask the " + "user to sign in again.", + config.org_token_endpoint, + err.summary, + ) + return assertion + case Ok(token): + verbose_proxy_logger.info( + "ID-JAG: exchanged the login assertion for a refresh token at %s (expires_in=%s)", + config.org_token_endpoint, + token.expires_in, + ) + return with_login_refresh_token( + assertion, + token.access_token, + token.expires_in, + now, + endpoint=config.org_token_endpoint, + client_id=config.client_id, + ) + + +async def exchange_login_assertion_if_configured( + assertion: SSOIdentityAssertion | None, +) -> SSOIdentityAssertion | None: + """The SSO-callback hook: consults the configured and DB-declared MCP servers and performs the + exchange when one of them presents a login refresh token. Nothing here can fail the login.""" + if assertion is None: + return None + from litellm.proxy._experimental.mcp_server.db import get_all_mcp_servers # noqa: PLC0415 # import cycle + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # avoids import cycle + global_mcp_server_manager, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( # noqa: PLC0415 # import cycle + to_server_spec, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.token_endpoint import ( # noqa: PLC0415 # cycle + TokenEndpointClient, + ) + from litellm.proxy.proxy_server import prisma_client # noqa: PLC0415 # runtime global + + try: + db_rows: Final = () if prisma_client is None else await get_all_mcp_servers(prisma_client) + servers: Final = await login_exchange_servers(global_mcp_server_manager, db_rows) + targets: Final = refresh_token_exchange_targets(to_server_spec(server) for server in servers) + return await exchange_login_assertion_for_refresh_token( + assertion, targets, TokenEndpointClient(), datetime.now(timezone.utc) + ) + except Exception as exc: # noqa: BLE001 # the login itself must not fail on an egress-side exchange + verbose_proxy_logger.warning("ID-JAG: the login-time refresh token exchange failed unexpectedly: %s", exc) + return assertion diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py index 66b44e91fd7..bd0813dd966 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py @@ -20,6 +20,7 @@ v2: no imports from v1. from __future__ import annotations import hashlib +from dataclasses import dataclass from datetime import datetime, timezone from functools import partial from typing import Final @@ -53,7 +54,10 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_r from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( AssertionStoreUnavailable, SSOAssertionStore, + SSOIdentityAssertion, assertion_expired, + exchanged_refresh_token_issued_for, + login_refresh_token_expired, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.token_endpoint import ( ExchangedToken, @@ -75,6 +79,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( ClientSecretAuth, CredError, IdJagConfig, + IdJagSubjectSource, NoneConfig, PassthroughConfig, PrivateKeyJwtAuth, @@ -88,6 +93,15 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( _TOKEN_EXCHANGE_GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:token-exchange" _JWT_BEARER_GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:jwt-bearer" _ID_JAG_REQUESTED_TOKEN_TYPE: Final = "urn:ietf:params:oauth:token-type:id-jag" +_REFRESH_TOKEN_SUBJECT_TOKEN_TYPE: Final = "urn:ietf:params:oauth:token-type:refresh_token" + + +@dataclass(frozen=True, slots=True) +class _IdJagSubjectToken: + """What ID-JAG leg 1 asserts and the RFC 8693 type it is declared as.""" + + value: str + token_type: str class _NullOAuthTokenStore: @@ -200,13 +214,15 @@ class UpstreamCredentialProvider: async def _id_jag( self, subject: Subject, server: ServerSpec, config: IdJagConfig ) -> Result[httpx2.Auth, CredError]: - match await self._id_jag_subject_token(subject): + match await self._id_jag_subject_token(subject, config): case Error(err): return Error(err) case Ok(subject_token): return await self._id_jag_exchange(subject, subject_token, server, config) - async def _id_jag_subject_token(self, subject: Subject) -> Result[str, CredError]: + async def _id_jag_subject_token( + self, subject: Subject, config: IdJagConfig + ) -> Result[_IdJagSubjectToken, CredError]: """The identity token ID-JAG leg 1 asserts, from the request or from the SSO login it was captured at. A caller that presents its own IdP identity token wins: that is the strongest available @@ -223,7 +239,7 @@ class UpstreamCredentialProvider: a raised error so a DB blip cannot 500 the egress or the upstream-401 retry. """ if subject.inbound_token is not None: - return Ok(subject.inbound_token.get_secret_value()) + return Ok(_IdJagSubjectToken(subject.inbound_token.get_secret_value(), config.subject_token_type)) if not subject.subject_id: return Error( CredError.of_precondition_required( @@ -232,7 +248,11 @@ class UpstreamCredentialProvider: ) ) try: - assertion: Final = await self._sso_assertion_store.fetch(subject.subject_id) + assertion: Final = await ( + self._sso_assertion_store.fetch_stored(subject.subject_id) + if config.subject_source is IdJagSubjectSource.refresh_token + else self._sso_assertion_store.fetch(subject.subject_id) + ) except AssertionStoreUnavailable as exc: # The driver's message can name hosts, schemas or connection details, and this summary # is returned to the caller verbatim as a 503 body. Operators get it from the log. @@ -253,6 +273,8 @@ class UpstreamCredentialProvider: "Sign in through LiteLLM SSO so the gateway captures one." ) ) + if config.subject_source is IdJagSubjectSource.refresh_token: + return _login_refresh_token_subject(assertion, config, datetime.now(timezone.utc)) if assertion_expired(assertion, datetime.now(timezone.utc)): return Error( CredError.of_precondition_required( @@ -260,20 +282,24 @@ class UpstreamCredentialProvider: "LiteLLM SSO again to capture a current one." ) ) - return Ok(assertion.id_token.get_secret_value()) + return Ok( + _IdJagSubjectToken( + assertion.id_token.get_secret_value(), assertion.subject_token_type or config.subject_token_type + ) + ) async def _id_jag_exchange( - self, subject: Subject, token: str, server: ServerSpec, config: IdJagConfig + self, subject: Subject, token: _IdJagSubjectToken, server: ServerSpec, config: IdJagConfig ) -> Result[httpx2.Auth, CredError]: slot: Final = _id_jag_slot_key(subject, server) - fingerprint: Final = _id_jag_fingerprint(token, server.server_id, config) + fingerprint: Final = _id_jag_fingerprint(token.value, server.server_id, config) async def _exchange() -> Result[ExchangedToken, CredError]: leg1_params: Final = { "grant_type": _TOKEN_EXCHANGE_GRANT_TYPE, "requested_token_type": _ID_JAG_REQUESTED_TOKEN_TYPE, - "subject_token": token, - "subject_token_type": config.subject_token_type, + "subject_token": token.value, + "subject_token_type": token.token_type, **({"audience": config.audience} if config.audience else {}), **({"resource": config.resource} if config.resource else {}), **({"scope": " ".join(config.scopes)} if config.scopes else {}), @@ -399,6 +425,41 @@ class UpstreamCredentialProvider: return None +def _login_refresh_token_subject( + assertion: SSOIdentityAssertion, config: IdJagConfig, now: datetime +) -> Result[_IdJagSubjectToken, CredError]: + """The subject for a server configured with ``id_jag_subject_source: refresh_token``: the refresh + token the ID-JAG client exchanged the login assertion for when the user signed in. Never the SSO + refresh token an OIDC login holds, that one is bound to the SSO client. The assertion's own expiry + does not apply here, a SAML assertion is spent minutes after login while the refresh token it + bought lives as long as the IdP said.""" + if assertion.exchanged_refresh_token is None: + return Error( + CredError.of_precondition_required( + "This server mints its ID-JAG from a refresh token exchanged at SSO login and none is stored " + "for this user. Sign in through LiteLLM SSO again; if the login log shows the IdP refused the " + "exchange, the IdP client has to allow a refresh token for this requesting app." + ) + ) + if not exchanged_refresh_token_issued_for(assertion, config.org_token_endpoint, config.client_id): + return Error( + CredError.of_precondition_required( + "The refresh token exchanged at SSO login was minted for a different IdP client than this server " + "uses, so it cannot be presented here. Point every id_jag_subject_source: refresh_token server at " + "the same IdP client, then sign in through LiteLLM SSO again." + ) + ) + if login_refresh_token_expired(assertion, now): + return Error( + CredError.of_precondition_required( + "The refresh token exchanged at SSO login for this user has expired. Sign in through LiteLLM SSO again." + ) + ) + return Ok( + _IdJagSubjectToken(assertion.exchanged_refresh_token.get_secret_value(), _REFRESH_TOKEN_SUBJECT_TOKEN_TYPE) + ) + + def _id_jag_slot_key(subject: Subject, server: ServerSpec) -> str: """Which cache slot this caller's bearer for this upstream lives in. diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_refresher.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_refresher.py index c41730ce9ec..ae5e9d0417e 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_refresher.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_refresher.py @@ -80,7 +80,9 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_s SSOIdentityAssertion, assertion_expired, assertion_from_sso_login, + carry_exchanged_refresh_token, fetch_sso_identity_assertion, + is_saml_assertion, persist_sso_identity_assertion, ) from litellm.types.llms.custom_http import httpxSpecialProvider @@ -293,6 +295,13 @@ class SSOAssertionRefresher: async def refresh( self, user_id: str, assertion: SSOIdentityAssertion ) -> Result[SSOIdentityAssertion, RefreshFailure]: + if is_saml_assertion(assertion): + verbose_proxy_logger.warning( + "ID-JAG: the stored SAML assertion for user_id=%s has reached its NotOnOrAfter. A SAML assertion " + "cannot be renewed, so the user has to sign in through SAML SSO again.", + user_id, + ) + return Error(RefreshFailure.of_rejected("a SAML assertion cannot be renewed; sign in again")) if assertion.refresh_token is None: verbose_proxy_logger.warning( "ID-JAG: the stored IdP identity assertion for user_id=%s has expired and no refresh token was " @@ -358,10 +367,11 @@ class SSOAssertionRefresher: user_id, ) return Error(RefreshFailure.of_rejected("the IdP's refresh response carried no usable id_token")) - failure: Final = await self._store_renewal(user_id, previous, renewed) + carried: Final = carry_exchanged_refresh_token(previous, renewed) + failure: Final = await self._store_renewal(user_id, previous, carried) if failure is not None: return Error(failure) - return Ok(renewed) + return Ok(carried) async def _store_renewal( self, user_id: str, previous: SSOIdentityAssertion, renewed: SSOIdentityAssertion @@ -436,9 +446,12 @@ class RefreshingSSOAssertionStore: self._skew = timedelta(seconds=expiry_skew_seconds) self._clock = clock + async def fetch_stored(self, user_id: str) -> SSOIdentityAssertion | None: + return await self._inner.fetch(user_id) + async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: assertion: Final = await self._inner.fetch(user_id) - if not self._expiring(assertion): + if assertion is None or is_saml_assertion(assertion) or not self._expiring(assertion): return assertion await self._coordinator().run( user_id, diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py index ef30b3d75b6..58494654082 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py @@ -1,9 +1,9 @@ """Store for the enterprise IdP identity assertion captured at SSO login (EMA). -The ``oauth2_id_jag`` egress arm needs the user's IdP ``id_token`` as its RFC 8693 -``subject_token``. A front-door client holds an identity-only ``llm_session_`` bearer, not an -IdP assertion, so the assertion captured at the one SSO login is the only usable subject -source for it. This module owns both sides of that state: the SSO callback persists here +The ``oauth2_id_jag`` egress arm needs the user's IdP identity assertion as its RFC 8693 +``subject_token``: the OIDC ``id_token``, or the signed SAML assertion when the login was SAML. +A front-door client holds an identity-only ``llm_session_`` bearer, not an IdP assertion, so +the assertion captured at the one SSO login is the only usable subject source for it. This module owns both sides of that state: the SSO callback persists here (write-through to the DB so a login on one pod is visible to every pod) and the resolver seam reads back by ``user_id``. Retention is gated on an ``oauth2_id_jag`` server actually being registered, so a gateway with no EMA upstream never stores bearer material. @@ -17,9 +17,10 @@ TTL ``MCP_SSO_ASSERTION_CACHE_TTL_SECONDS``; invalidation also guards against st from __future__ import annotations +import base64 import json from collections.abc import Mapping, Sequence -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from types import MappingProxyType from typing import TYPE_CHECKING, Final, Protocol @@ -44,6 +45,7 @@ GOOGLE_REFRESH_TOKEN_PARAMS: Final = MappingProxyType({"access_type": "offline", _ASSERTION_DECRYPT_LOG_KEY: Final = "sso_identity_assertion" +SAML2_SUBJECT_TOKEN_TYPE: Final = "urn:ietf:params:oauth:token-type:saml2" _STR_ADAPTER: Final[TypeAdapter[str]] = TypeAdapter(str) _MAYBE_STR_ADAPTER: Final[TypeAdapter[str | None]] = TypeAdapter(str | None) @@ -77,8 +79,10 @@ def _mcp_server_table(prisma_client: PrismaClient) -> _MCPServerTable: class SSOIdentityAssertion(LiteLLMBaseModel): - """The IdP material an EMA exchange needs: ``id_token`` is the RFC 8693 subject token, - ``expires_at`` bounds its usefulness, and the refresh token renews it without re-login.""" + """The IdP material an EMA exchange needs: ``id_token`` is the RFC 8693 subject token (the OIDC + id_token, or the base64 signed ```` when ``subject_token_type`` names SAML), + ``expires_at`` bounds its usefulness, and the refresh token renews it without re-login. + ``subject_token_type`` is ``None`` when the exchange should send the operator-configured type.""" model_config = ConfigDict(frozen=True) @@ -86,6 +90,11 @@ class SSOIdentityAssertion(LiteLLMBaseModel): refresh_token: SecretStr | None = None issuer: str | None = None expires_at: datetime | None = None + subject_token_type: str | None = None + exchanged_refresh_token: SecretStr | None = None + exchanged_refresh_token_expires_at: datetime | None = None + exchanged_refresh_token_endpoint: str | None = None + exchanged_refresh_token_client_id: str | None = None class SSOAssertionCache: @@ -138,6 +147,11 @@ class _StoredAssertionPayload(LiteLLMBaseModel): refresh_token: str | None = None issuer: str | None = None expires_at: datetime | None = None + subject_token_type: str | None = None + exchanged_refresh_token: str | None = None + exchanged_refresh_token_expires_at: datetime | None = None + exchanged_refresh_token_endpoint: str | None = None + exchanged_refresh_token_client_id: str | None = None def assertion_from_sso_login(id_token: object, refresh_token: object) -> SSOIdentityAssertion | None: @@ -167,6 +181,85 @@ def assertion_from_sso_login(id_token: object, refresh_token: object) -> SSOIden ) +def assertion_from_saml_login( + assertion_xml: bytes, not_on_or_after: int | None, issuer: str | None +) -> SSOIdentityAssertion: + """The carrier for a SAML login: the signed ```` element the IdP issued, base64 + encoded the way RFC 8693 carries a ``saml2`` subject token, with its own validity window as + the expiry and no refresh token, since SAML has none and the user has to sign in again.""" + return SSOIdentityAssertion( + id_token=SecretStr(base64.urlsafe_b64encode(assertion_xml).rstrip(b"=").decode("ascii")), + issuer=issuer, + expires_at=datetime.fromtimestamp(not_on_or_after, tz=timezone.utc) if not_on_or_after is not None else None, + subject_token_type=SAML2_SUBJECT_TOKEN_TYPE, + ) + + +def is_saml_assertion(assertion: SSOIdentityAssertion) -> bool: + return assertion.subject_token_type == SAML2_SUBJECT_TOKEN_TYPE + + +def with_login_refresh_token( + assertion: SSOIdentityAssertion, + refresh_token: str, + expires_in: int | None, + now: datetime, + *, + endpoint: str | None = None, + client_id: str | None = None, +) -> SSOIdentityAssertion: + """The same login, now also carrying the refresh token the ID-JAG client exchanged the assertion + for, with the lifetime the IdP gave that token. It is kept apart from ``refresh_token``, which an + OIDC login holds for the SSO client and which only that client can redeem. The assertion and its + own expiry stay as they were, so a server that presents the login assertion directly is unaffected.""" + return assertion.model_copy( + update={ + "exchanged_refresh_token": SecretStr(refresh_token), + "exchanged_refresh_token_expires_at": now + timedelta(seconds=expires_in) + if expires_in is not None + else None, + "exchanged_refresh_token_endpoint": endpoint, + "exchanged_refresh_token_client_id": client_id, + } + ) + + +def carry_exchanged_refresh_token( + previous: SSOIdentityAssertion, renewed: SSOIdentityAssertion +) -> SSOIdentityAssertion: + """The renewed OIDC assertion still carrying the refresh token the ID-JAG client exchanged at login. + Renewal only replaces the id_token and the SSO refresh token, the exchanged one is unrelated to + them and stays valid for as long as the IdP said.""" + return renewed.model_copy( + update={ + "exchanged_refresh_token": previous.exchanged_refresh_token, + "exchanged_refresh_token_expires_at": previous.exchanged_refresh_token_expires_at, + "exchanged_refresh_token_endpoint": previous.exchanged_refresh_token_endpoint, + "exchanged_refresh_token_client_id": previous.exchanged_refresh_token_client_id, + } + ) + + +def exchanged_refresh_token_issued_for(assertion: SSOIdentityAssertion, endpoint: str, client_id: str) -> bool: + """Whether the exchanged refresh token was minted at this token endpoint for this client. A refresh + token is bound to the client that received it, so one minted for another server's IdP client must + never be posted to this server's endpoint.""" + return ( + assertion.exchanged_refresh_token_endpoint == endpoint + and assertion.exchanged_refresh_token_client_id == client_id + ) + + +def login_refresh_token_expired(assertion: SSOIdentityAssertion, now: datetime) -> bool: + """Whether the refresh token exchanged at login has passed the lifetime the IdP gave it. One + with no recorded lifetime is left for the IdP to judge, like an assertion with no ``exp``.""" + expires_at: Final = assertion.exchanged_refresh_token_expires_at + if expires_at is None: + return False + normalized: Final = expires_at if expires_at.tzinfo is not None else expires_at.replace(tzinfo=timezone.utc) + return normalized <= now + + def assertion_expired(assertion: SSOIdentityAssertion, now: datetime) -> bool: """Whether the assertion's ``exp`` has passed at ``now``. An assertion carrying no expiry is treated as usable and left for the IdP to reject, since the store records what the id_token @@ -220,6 +313,27 @@ async def persist_sso_identity_assertion( **({"refresh_token": assertion.refresh_token.get_secret_value()} if assertion.refresh_token else {}), **({"issuer": assertion.issuer} if assertion.issuer else {}), **({"expires_at": assertion.expires_at.isoformat()} if assertion.expires_at else {}), + **({"subject_token_type": assertion.subject_token_type} if assertion.subject_token_type else {}), + **( + {"exchanged_refresh_token": assertion.exchanged_refresh_token.get_secret_value()} + if assertion.exchanged_refresh_token + else {} + ), + **( + {"exchanged_refresh_token_expires_at": assertion.exchanged_refresh_token_expires_at.isoformat()} + if assertion.exchanged_refresh_token_expires_at + else {} + ), + **( + {"exchanged_refresh_token_endpoint": assertion.exchanged_refresh_token_endpoint} + if assertion.exchanged_refresh_token_endpoint + else {} + ), + **( + {"exchanged_refresh_token_client_id": assertion.exchanged_refresh_token_client_id} + if assertion.exchanged_refresh_token_client_id + else {} + ), } encoded: Final = _STR_ADAPTER.validate_python(encrypt_value_helper(json.dumps(payload))) await _assertion_table(prisma_client).upsert( @@ -258,6 +372,11 @@ async def _read_assertion_from_db(user_id: str) -> SSOIdentityAssertion | None: refresh_token=SecretStr(payload.refresh_token) if payload.refresh_token else None, issuer=payload.issuer, expires_at=payload.expires_at, + subject_token_type=payload.subject_token_type, + exchanged_refresh_token=SecretStr(payload.exchanged_refresh_token) if payload.exchanged_refresh_token else None, + exchanged_refresh_token_expires_at=payload.exchanged_refresh_token_expires_at, + exchanged_refresh_token_endpoint=payload.exchanged_refresh_token_endpoint, + exchanged_refresh_token_client_id=payload.exchanged_refresh_token_client_id, ) @@ -298,6 +417,12 @@ class SSOAssertionStore(Protocol): async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: ... + async def fetch_stored(self, user_id: str) -> SSOIdentityAssertion | None: + """The row as stored, never renewed on the way out. For a server that presents the refresh + token exchanged at login, the SSO id_token's own lifetime is irrelevant and an SSO renewal + failing must not stand between that server and a still valid refresh token.""" + ... + class DbSSOAssertionStore: """The live store: the row the SSO callback wrote, read back by ``user_id``. @@ -316,6 +441,9 @@ class DbSSOAssertionStore: except Exception as exc: # noqa: BLE001 # any driver/storage failure is an outage, not an absence raise AssertionStoreUnavailable(str(exc)) from exc + async def fetch_stored(self, user_id: str) -> SSOIdentityAssertion | None: + return await self.fetch(user_id) + async def fetch_uncached(self, user_id: str) -> SSOIdentityAssertion | None: try: return await _read_assertion_from_db(user_id) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py index 71745a560f3..45c53af9349 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py @@ -349,6 +349,14 @@ class IdJagResourceClient(LiteLLMBaseModel): client_auth: ClientAuth +class IdJagSubjectSource(str, Enum): + """What ID-JAG leg 1 presents as the RFC 8693 subject_token: the identity assertion captured at SSO + login, or the refresh token that assertion was exchanged for at the IdP when the user logged in.""" + + login_assertion = "login_assertion" + refresh_token = "refresh_token" + + class IdJagConfig(HeaderCarrier): """draft-ietf-oauth-identity-assertion-authz-grant (Okta "AI agent token exchange"). @@ -368,6 +376,7 @@ class IdJagConfig(HeaderCarrier): client_auth: ClientAuth resource_client: IdJagResourceClient subject_token_type: str = "urn:ietf:params:oauth:token-type:id_token" + subject_source: IdJagSubjectSource = IdJagSubjectSource.login_assertion audience: str | None = None resource: str | None = None scopes: tuple[str, ...] = () diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 2eb62a07447..467cb2d624c 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -34541,6 +34541,17 @@ ], "title": "Id Jag Resource Token Endpoint" }, + "id_jag_subject_source": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Id Jag Subject Source" + }, "redirect_uris": { "anyOf": [ { @@ -37358,6 +37369,17 @@ ], "title": "Id Jag Resource Token Endpoint" }, + "id_jag_subject_source": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Id Jag Subject Source" + }, "redirect_uris": { "anyOf": [ { @@ -40763,6 +40785,17 @@ ], "title": "Id Jag Resource Token Endpoint" }, + "id_jag_subject_source": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Id Jag Subject Source" + }, "redirect_uris": { "anyOf": [ { @@ -44843,6 +44876,17 @@ ], "title": "Id Jag Resource Token Endpoint" }, + "id_jag_subject_source": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Id Jag Subject Source" + }, "redirect_uris": { "anyOf": [ { diff --git a/litellm/proxy/management_endpoints/sso/id_jag_assertion_capture.py b/litellm/proxy/management_endpoints/sso/id_jag_assertion_capture.py index deb9d2975de..ff7e68c3ab2 100644 --- a/litellm/proxy/management_endpoints/sso/id_jag_assertion_capture.py +++ b/litellm/proxy/management_endpoints/sso/id_jag_assertion_capture.py @@ -2,10 +2,10 @@ An ``oauth2_id_jag`` MCP server spends the ``id_token`` captured at SSO login as its RFC 8693 subject token. The Google, Microsoft and generic OIDC login paths all reach a token response the -gateway retains one from; SAML has no id_token to capture, and a deployment with no SSO never signs -anyone in, so under either every store-sourced ID-JAG exchange fails for every user, however many -times they sign in. Neither side can see that alone: the MCP registration knows nothing about SSO -and the login knows nothing about MCP. This module is the one shared answer both warn from. +gateway retains one from, and the SAML login retains the signed assertion itself as a ``saml2`` +subject token. A deployment with no SSO never signs anyone in, so every store-sourced ID-JAG +exchange fails for every user. Neither side can see that alone: the MCP registration knows nothing +about SSO and the login knows nothing about MCP. This module is the one shared answer both warn from. """ from __future__ import annotations @@ -19,8 +19,8 @@ from typing_extensions import assert_never from litellm.proxy.management_endpoints.sso.saml_sso import SAMLAuthHandler _OIDC_REMEDY: Final = ( - "Point SSO at an OIDC provider (GOOGLE_CLIENT_ID, MICROSOFT_CLIENT_ID or GENERIC_CLIENT_ID), the login " - "paths whose token response the gateway retains an id_token from" + "Point SSO at an OIDC provider (GOOGLE_CLIENT_ID, MICROSOFT_CLIENT_ID or GENERIC_CLIENT_ID) or at SAML, " + "the login paths the gateway retains an identity assertion from" ) @@ -56,19 +56,15 @@ def id_jag_assertion_capture_gap() -> str | None: or ``None`` when that provider does capture an assertion.""" provider = active_sso_provider() match provider: - case ActiveSSOProvider.google | ActiveSSOProvider.microsoft | ActiveSSOProvider.generic: + case ( + ActiveSSOProvider.google | ActiveSSOProvider.microsoft | ActiveSSOProvider.generic | ActiveSSOProvider.saml + ): return None case ActiveSSOProvider.none: return ( "no SSO provider is configured, so no IdP identity assertion is ever captured and " f"ID-JAG credential resolution fails for every user. {_OIDC_REMEDY}" ) - case ActiveSSOProvider.saml: - return ( - f"the active SSO provider ({provider.value}) has no identity-assertion capture path, so no " - "IdP id_token is ever stored and ID-JAG credential resolution fails for every user no matter " - f"how often they sign in. {_OIDC_REMEDY}" - ) case _: assert_never(provider) diff --git a/litellm/proxy/management_endpoints/sso/saml_sso.py b/litellm/proxy/management_endpoints/sso/saml_sso.py index 4991436edac..ba265acfe4a 100644 --- a/litellm/proxy/management_endpoints/sso/saml_sso.py +++ b/litellm/proxy/management_endpoints/sso/saml_sso.py @@ -27,6 +27,7 @@ import hashlib import os import secrets import time +from collections.abc import Mapping from typing import Final, cast from urllib.parse import parse_qsl @@ -37,6 +38,10 @@ from pydantic import ValidationError from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.caching.dual_cache import DualCache +from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( + SSOIdentityAssertion, + assertion_from_saml_login, +) from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.proxy.management_endpoints.sso_helper_utils import SSO_SESSIONS_TARGET from litellm.proxy.management_endpoints.types import CustomOpenID, get_litellm_user_role @@ -212,7 +217,7 @@ class SAMLAuthHandler: return OneLogin_Saml2_IdPMetadataParser.merge_settings(sp_settings, idp_settings) @staticmethod - def _prepare_request_data(request: Request, post_data: dict[str, str] | None = None) -> dict[str, object]: + def _prepare_request_data(request: Request, post_data: Mapping[str, str] | None = None) -> dict[str, object]: base: Final = SAMLAuthHandler._base_url(request) scheme, _, host_part = base.partition("://") host: Final = host_part.split("/", 1)[0] @@ -228,7 +233,7 @@ class SAMLAuthHandler: async def _build_auth( request: Request, cache: DualCache, - post_data: dict[str, str] | None = None, + post_data: Mapping[str, str] | None = None, ) -> "OneLogin_Saml2_Auth": if not SAML_AVAILABLE: raise _saml_unavailable_error() @@ -310,7 +315,9 @@ class SAMLAuthHandler: return dict(parse_qsl(body.decode("utf-8", "replace"))) @staticmethod - async def handle_acs(request: Request, cache: DualCache, post_data: dict[str, str]) -> CustomOpenID: + async def handle_acs( + request: Request, cache: DualCache, post_data: Mapping[str, str] + ) -> tuple[CustomOpenID, SSOIdentityAssertion | None]: auth: Final = await SAMLAuthHandler._build_auth(request, cache, post_data=post_data) browser_request_id: Final = request.cookies.get(_SAML_AUTHN_STATE_COOKIE) try: @@ -330,7 +337,50 @@ class SAMLAuthHandler: ) await SAMLAuthHandler._enforce_response_binding(auth, cache, browser_request_id) - return SAMLAuthHandler._result_from_auth(auth) + return SAMLAuthHandler._result_from_auth(auth), SAMLAuthHandler._retained_assertion(auth) + + @staticmethod + def _retained_assertion(auth: "OneLogin_Saml2_Auth") -> SSOIdentityAssertion | None: + """The one signed ```` the validated response carried, serialized on its own + so its enveloped signature still verifies at the IdP's token endpoint. Retention is optional + egress material, so a failure here is logged and the validated login goes on without it.""" + try: + return SAMLAuthHandler._parse_retained_assertion(auth) + except Exception as e: # noqa: BLE001 - lxml and python3-saml expose no common base; capture must not block login + verbose_proxy_logger.warning( + "SAML login validated but its assertion could not be retained for oauth2_id_jag MCP servers: %s", e + ) + return None + + @staticmethod + def _parse_retained_assertion(auth: "OneLogin_Saml2_Auth") -> SSOIdentityAssertion | None: + xml: Final = cast(bytes | None, auth.get_last_response_xml()) # cast-ok: untyped python3-saml + if not xml: + verbose_proxy_logger.warning( + "SAML login validated but the response XML was not retained, so this user has no identity " + "assertion for oauth2_id_jag MCP servers." + ) + return None + root: Final = OneLogin_Saml2_XML.to_etree(xml) + assertions: Final = OneLogin_Saml2_XML.query(root, "//saml:Assertion") + if len(assertions) != 1: + encrypted: Final = OneLogin_Saml2_XML.query(root, "//saml:EncryptedAssertion") + verbose_proxy_logger.warning( + "SAML login validated but the response carried %d plaintext saml:Assertion elements and %d " + "saml:EncryptedAssertion elements, so no identity assertion was retained for oauth2_id_jag MCP " + "servers. Those servers need exactly one signed, unencrypted assertion from the IdP.", + len(assertions), + len(encrypted), + ) + return None + issuers: Final = OneLogin_Saml2_XML.query(assertions[0], "./saml:Issuer") + issuer: Final = cast(str | None, issuers[0].text) if issuers else None # cast-ok: untyped python3-saml + not_on_or_after: Final = auth.get_last_assertion_not_on_or_after() + return assertion_from_saml_login( + assertion_xml=cast(bytes, OneLogin_Saml2_XML.to_string(assertions[0], encoding="utf-8")), # cast-ok: saml + not_on_or_after=not_on_or_after if isinstance(not_on_or_after, int) else None, + issuer=issuer, + ) @staticmethod def _replay_guard_ttl(auth: "OneLogin_Saml2_Auth") -> int: diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index c3c5d7401c5..e6d93aa6260 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -71,6 +71,9 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.login_assertion_exchange import ( + exchange_login_assertion_if_configured, +) from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( GOOGLE_REFRESH_TOKEN_PARAMS, MICROSOFT_REFRESH_TOKEN_SCOPE, @@ -2343,7 +2346,7 @@ async def saml_callback(request: Request): if "SAMLResponse" not in post_data: raise HTTPException(status_code=400, detail="Missing SAMLResponse in callback request.") - result: Final = await SAMLAuthHandler.handle_acs(request=request, cache=saml_sso_cache, post_data=post_data) + user, retained = await SAMLAuthHandler.handle_acs(request=request, cache=saml_sso_cache, post_data=post_data) await _raise_if_sso_exceeds_free_user_limit(premium_user, prisma_client) @@ -2356,7 +2359,7 @@ async def saml_callback(request: Request): ) return await SSOAuthenticationHandler.get_redirect_response_from_openid( - result=result, + result=user, request=request, received_response=None, generic_client_id=None, @@ -2364,6 +2367,7 @@ async def saml_callback(request: Request): access_token_payload=None, jwt_handler=jwt_handler, return_to=cp_return_to, + sso_assertion=retained, ) @@ -2480,7 +2484,8 @@ async def _complete_cli_sso_callback_session( raise HTTPException(status_code=500, detail="Failed to retrieve user information from SSO") resolved_user_id: Final = _require_sso_user_id(user_info.user_id) - await retain_sso_identity_assertion_for_ema(user_id=resolved_user_id, assertion=sso_assertion) + retained_assertion: Final = await exchange_login_assertion_if_configured(sso_assertion) + await retain_sso_identity_assertion_for_ema(user_id=resolved_user_id, assertion=retained_assertion) await warn_if_id_jag_assertion_uncaptured(sso_assertion) teams: list[str] = [] @@ -3863,7 +3868,8 @@ class SSOAuthenticationHandler: ) if isinstance(user_id, str) and user_id: - await retain_sso_identity_assertion_for_ema(user_id=user_id, assertion=sso_assertion) + retained_assertion: Final = await exchange_login_assertion_if_configured(sso_assertion) + await retain_sso_identity_assertion_for_ema(user_id=user_id, assertion=retained_assertion) await warn_if_id_jag_assertion_uncaptured(sso_assertion) disabled_non_admin_personal_key_creation: Final = get_disabled_non_admin_personal_key_creation() diff --git a/litellm/types/mcp.py b/litellm/types/mcp.py index 72e564b6919..f0be66f055d 100644 --- a/litellm/types/mcp.py +++ b/litellm/types/mcp.py @@ -248,6 +248,15 @@ class MCPCredentials(TypedDict, total=False): Optional RFC 8707 resource indicator sent on ID-JAG leg 1 """ + id_jag_subject_source: ReadOnly[str | None] + """ + Which token ID-JAG leg 1 presents as its RFC 8693 subject_token. ``login_assertion`` (the default) is the + identity assertion captured at SSO login, an OIDC id_token or a SAML assertion. ``refresh_token`` exchanges + that assertion at the IdP token endpoint once, at login, for a refresh token and presents the refresh token + on every leg 1 instead, which is how a SAML requesting app works at IdPs whose ID-JAG exchange does not take + a saml2 subject directly. + """ + id_jag_resource_client_id: ReadOnly[str | None] """OAuth client identifier used at the ID-JAG resource authorization server.""" diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index f7c4d329781..219f23200d2 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -182,6 +182,7 @@ class MCPServer(LiteLLMBaseModel): # posts the ID-JAG assertion to id_jag_resource_token_endpoint. id_jag_resource_token_endpoint: str | None = None id_jag_resource: str | None = None + id_jag_subject_source: str | None = None id_jag_resource_client_id: str | None = Field( default=None, description="OAuth client identifier used at the ID-JAG resource authorization server.", diff --git a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py index 2cdf716fa30..9cd19e9a1c9 100644 --- a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py +++ b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py @@ -31,6 +31,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( ClientSecretAuth, CredError, IdJagConfig, + IdJagSubjectSource, NoneConfig, PassthroughConfig, PrivateKeyJwtAuth, @@ -908,3 +909,25 @@ def test_incomplete_obo_keeps_exchange_ownership(client_secret: str | None, is_b assert isinstance(spec.config, TokenExchangeConfig) assert spec.config.client_id == "client" assert spec.config.client_secret is None + assert spec.config.client_id == "client" + assert spec.config.client_secret is None + + +@pytest.mark.parametrize( + ("configured", "expected"), + [ + (None, IdJagSubjectSource.login_assertion), + ("login_assertion", IdJagSubjectSource.login_assertion), + ("refresh_token", IdJagSubjectSource.refresh_token), + ], +) +def test_id_jag_subject_source_maps_to_config(configured: str | None, expected: IdJagSubjectSource) -> None: + spec = to_server_spec(_id_jag_server(id_jag_subject_source=configured)) + + assert spec is not None + assert isinstance(spec.config, IdJagConfig) + assert spec.config.subject_source is expected + + +def test_id_jag_unknown_subject_source_defers_like_a_half_configured_server() -> None: + assert to_server_spec(_id_jag_server(id_jag_subject_source="assertion")) is None diff --git a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_login_assertion_exchange.py b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_login_assertion_exchange.py new file mode 100644 index 00000000000..bb228544e14 --- /dev/null +++ b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_login_assertion_exchange.py @@ -0,0 +1,327 @@ +"""The login-time exchange of a captured identity assertion for an IdP refresh token.""" + +import asyncio +import logging +from datetime import datetime, timedelta, timezone + +import pytest +from pydantic import SecretStr + +from litellm.proxy._experimental.mcp_server.outbound_credentials.login_assertion_exchange import ( + exchange_login_assertion_for_refresh_token, + login_exchange_servers, + refresh_token_exchange_targets, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok, Result +from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( + SAML2_SUBJECT_TOKEN_TYPE, + SSOIdentityAssertion, + assertion_from_saml_login, + login_refresh_token_expired, + with_login_refresh_token, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.token_endpoint import ExchangedToken +from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( + ClientAuth, + ClientSecretAuth, + CredError, + IdJagConfig, + IdJagResourceClient, + IdJagSubjectSource, + NoneConfig, + ServerSpec, +) +from litellm.proxy._types import LiteLLM_MCPServerTable +from litellm.types.mcp import MCPAuth, MCPTransport +from litellm.types.mcp_server.mcp_server_manager import MCPServer + +_NOW = datetime(2026, 10, 16, 12, 0, tzinfo=timezone.utc) + + +def _config( + *, source: IdJagSubjectSource = IdJagSubjectSource.refresh_token, client_id: str = "agent", endpoint: str = "" +) -> IdJagConfig: + return IdJagConfig( + org_token_endpoint=endpoint or "https://idp.example.com/token", + resource_token_endpoint="https://mcp-as.example.com/token", + client_id=client_id, + client_auth=ClientSecretAuth(client_secret=SecretStr("s")), + resource_client=IdJagResourceClient( + client_id=client_id, client_auth=ClientSecretAuth(client_secret=SecretStr("s")) + ), + subject_source=source, + ) + + +def _spec(config: IdJagConfig | NoneConfig, server_id: str = "srv") -> ServerSpec: + return ServerSpec(server_id=server_id, resource="https://mcp.example.com", config=config) + + +def _saml() -> SSOIdentityAssertion: + return assertion_from_saml_login(b"", int(_NOW.timestamp()) + 300, "https://idp.example.com") + + +class _Endpoint: + def __init__(self, result: Result[ExchangedToken, CredError]) -> None: + self._result = result + self.calls: list[tuple[str, str, dict[str, str], ClientAuth]] = [] + + async def fetch( + self, endpoint: str, client_id: str, grant_params: dict[str, str], client_auth: ClientAuth + ) -> Result[ExchangedToken, CredError]: + self.calls.append((endpoint, client_id, dict(grant_params), client_auth)) + return self._result + + +def test_targets_keep_only_refresh_token_sourced_id_jag_servers_one_per_idp_client(): + direct = _config(source=IdJagSubjectSource.login_assertion) + first = _config(client_id="agent") + twin = _config(client_id="agent") + other = _config(client_id="other-agent") + + targets = refresh_token_exchange_targets( + [_spec(direct, "a"), None, _spec(NoneConfig(), "b"), _spec(first, "c"), _spec(other, "d"), _spec(twin, "e")] + ) + + assert targets == (first, other) + + +@pytest.mark.asyncio +async def test_exchange_presents_the_captured_assertion_and_keeps_the_refresh_token(): + endpoint = _Endpoint(Ok(ExchangedToken(access_token="rt-1", expires_in=3600))) + config = _config() + saml = _saml() + + exchanged = await exchange_login_assertion_for_refresh_token(saml, (config,), endpoint, _NOW) + + ((url, client_id, params, auth),) = endpoint.calls + assert (url, client_id, auth) == (config.org_token_endpoint, config.client_id, config.client_auth) + assert params == { + "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange", + "requested_token_type": "urn:ietf:params:oauth:token-type:refresh_token", + "subject_token": saml.id_token.get_secret_value(), + "subject_token_type": SAML2_SUBJECT_TOKEN_TYPE, + "scope": "openid offline_access", + } + assert exchanged.exchanged_refresh_token is not None + assert exchanged.exchanged_refresh_token.get_secret_value() == "rt-1" + assert exchanged.exchanged_refresh_token_endpoint == "https://idp.example.com/token" + assert exchanged.exchanged_refresh_token_client_id == "agent" + assert exchanged.exchanged_refresh_token_expires_at == _NOW + timedelta(seconds=3600) + assert exchanged.id_token == saml.id_token + assert exchanged.expires_at == saml.expires_at + assert exchanged.subject_token_type == SAML2_SUBJECT_TOKEN_TYPE + assert not login_refresh_token_expired(exchanged, _NOW + timedelta(seconds=3599)) + assert login_refresh_token_expired(exchanged, _NOW + timedelta(seconds=3600)) + + +@pytest.mark.asyncio +async def test_exchange_without_a_lifetime_leaves_expiry_to_the_idp(): + endpoint = _Endpoint(Ok(ExchangedToken(access_token="rt-1", expires_in=None))) + + exchanged = await exchange_login_assertion_for_refresh_token(_saml(), (_config(),), endpoint, _NOW) + + assert exchanged.exchanged_refresh_token_expires_at is None + assert not login_refresh_token_expired(exchanged, _NOW + timedelta(days=365)) + + +@pytest.mark.asyncio +async def test_a_refused_exchange_keeps_the_plain_assertion(caplog: pytest.LogCaptureFixture): + endpoint = _Endpoint(Error(CredError.of_upstream_unavailable("token exchange failed with status 400"))) + saml = _saml() + + with caplog.at_level("WARNING"): + exchanged = await exchange_login_assertion_for_refresh_token(saml, (_config(),), endpoint, _NOW) + + assert exchanged == saml + assert exchanged.exchanged_refresh_token is None + assert any("refused to exchange the login assertion" in message for message in caplog.messages) + + +@pytest.mark.asyncio +async def test_no_exchange_without_a_refresh_token_sourced_server(): + endpoint = _Endpoint(Ok(ExchangedToken(access_token="rt-1", expires_in=3600))) + saml = _saml() + + assert await exchange_login_assertion_for_refresh_token(saml, (), endpoint, _NOW) == saml + assert endpoint.calls == [] + + +@pytest.mark.asyncio +async def test_an_oidc_login_is_exchanged_at_the_id_jag_client_and_keeps_its_sso_refresh_token(): + """The SSO refresh token belongs to the SSO client and cannot be redeemed by the ID-JAG client, so + an OIDC login still exchanges its id_token and both tokens live side by side.""" + endpoint = _Endpoint(Ok(ExchangedToken(access_token="rt-1", expires_in=3600))) + oidc = SSOIdentityAssertion(id_token=SecretStr("id-token"), refresh_token=SecretStr("sso-rt")) + + exchanged = await exchange_login_assertion_for_refresh_token(oidc, (_config(),), endpoint, _NOW) + + assert endpoint.calls[0][2]["subject_token"] == "id-token" + assert endpoint.calls[0][2]["subject_token_type"] == "urn:ietf:params:oauth:token-type:id_token" + assert exchanged.refresh_token is not None and exchanged.refresh_token.get_secret_value() == "sso-rt" + assert exchanged.exchanged_refresh_token is not None + assert exchanged.exchanged_refresh_token.get_secret_value() == "rt-1" + + +@pytest.mark.asyncio +async def test_a_login_already_exchanged_is_not_exchanged_again(): + endpoint = _Endpoint(Ok(ExchangedToken(access_token="rt-2", expires_in=3600))) + done = with_login_refresh_token(_saml(), "rt-1", 3600, _NOW) + + assert await exchange_login_assertion_for_refresh_token(done, (_config(),), endpoint, _NOW) == done + assert endpoint.calls == [] + + +@pytest.mark.asyncio +async def test_several_idp_clients_mint_for_the_first_and_warn(caplog: pytest.LogCaptureFixture): + endpoint = _Endpoint(Ok(ExchangedToken(access_token="rt-1", expires_in=3600))) + first, second = _config(client_id="agent"), _config(client_id="other-agent") + + with caplog.at_level("WARNING"): + exchanged = await exchange_login_assertion_for_refresh_token(_saml(), (first, second), endpoint, _NOW) + + assert [call[1] for call in endpoint.calls] == ["agent"] + assert exchanged.exchanged_refresh_token is not None + assert any("2 IdP clients" in message and "client_id=agent" in message for message in caplog.messages) + + +class _FakeCatalog: + def __init__(self, config_servers: dict[str, MCPServer]) -> None: + self.config_mcp_servers = config_servers + self.built: list[tuple[str, bool]] = [] + + async def build_mcp_server_from_table( + self, + mcp_server: LiteLLM_MCPServerTable, + *, + credentials_are_encrypted: bool = True, + env_vars_are_encrypted: bool | None = None, + register_oauth_discovery: bool = True, + ) -> MCPServer: + self.built.append((mcp_server.server_id, register_oauth_discovery)) + return MCPServer( + server_id=mcp_server.server_id, name=mcp_server.server_id, url=mcp_server.url, transport=MCPTransport.http + ) + + +def _row(server_id: str, auth_type: MCPAuth, approval_status: str | None = None) -> LiteLLM_MCPServerTable: + return LiteLLM_MCPServerTable( + server_id=server_id, + alias=server_id, + url=f"https://{server_id}.example.com/mcp", + transport=MCPTransport.http, + auth_type=auth_type, + approval_status=approval_status, + created_at=_NOW, + updated_at=_NOW, + ) + + +@pytest.mark.asyncio +async def test_login_exchange_servers_reads_config_and_db_id_jag_rows_not_the_registry(): + """A server added on another pod is in the DB before this pod's registry reloads, so the login + exchange has to judge against the DB rows, like the retention gate, and only the ID-JAG ones.""" + config_server = MCPServer( + server_id="from-config", name="from-config", url="https://cfg.example.com/mcp", transport=MCPTransport.http + ) + catalog = _FakeCatalog({"from-config": config_server}) + + servers = await login_exchange_servers( + catalog, + [ + _row("db-id-jag", MCPAuth.oauth2_id_jag), + _row("db-id-jag-approved", MCPAuth.oauth2_id_jag, "approved"), + _row("db-oauth", MCPAuth.oauth2), + ], + ) + + assert [server.server_id for server in servers] == ["from-config", "db-id-jag", "db-id-jag-approved"] + assert catalog.built == [("db-id-jag", False), ("db-id-jag-approved", False)] + + +@pytest.mark.asyncio +async def test_a_db_row_hides_the_config_server_with_the_same_id_like_the_live_catalog(): + """The catalog serves ``config_mcp_servers | registry``, so a DB row wins its server_id. The login + exchange must mint for that live server, not for the hidden config entry's IdP client.""" + hidden = MCPServer( + server_id="shared", name="shared", url="https://cfg.example.com/mcp", transport=MCPTransport.http + ) + catalog = _FakeCatalog({"shared": hidden}) + + servers = await login_exchange_servers(catalog, [_row("shared", MCPAuth.oauth2_id_jag)]) + + assert [(server.server_id, server.url) for server in servers] == [("shared", "https://shared.example.com/mcp")] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("approval_status", ["pending_review", "rejected", "draft"]) +async def test_login_exchange_servers_skip_unapproved_submissions(approval_status: str): + """A team key can submit an ID-JAG server naming any token endpoint. Until an admin approves it, + no login assertion is posted there.""" + catalog = _FakeCatalog({}) + + servers = await login_exchange_servers(catalog, [_row("submitted", MCPAuth.oauth2_id_jag, approval_status)]) + + assert servers == () + assert catalog.built == [] + + +class _StalledEndpoint(_Endpoint): + """An IdP that accepts the connection and never answers.""" + + async def fetch( + self, endpoint: str, client_id: str, grant_params: dict[str, str], client_auth: ClientAuth + ) -> Result[ExchangedToken, CredError]: + await asyncio.sleep(0.5) + return await super().fetch(endpoint, client_id, grant_params, client_auth) + + +@pytest.mark.asyncio +async def test_a_stalled_idp_does_not_hold_up_the_login(caplog: pytest.LogCaptureFixture): + """The exchange is optional. A user who only wants the dashboard must not wait on an MCP IdP that + stopped answering, so the exchange has its own short deadline and the login keeps the plain assertion.""" + saml = _saml() + + with caplog.at_level(logging.WARNING, logger="litellm.proxy"): + exchanged = await exchange_login_assertion_for_refresh_token( + saml, + (_config(),), + _StalledEndpoint(Ok(ExchangedToken(access_token="rt-1", expires_in=3600))), + _NOW, + deadline_seconds=0.05, + ) + + assert exchanged == saml + assert any("did not answer the login-time refresh token exchange" in r.getMessage() for r in caplog.records) + + +class _OneBadRowCatalog(_FakeCatalog): + async def build_mcp_server_from_table( + self, + mcp_server: LiteLLM_MCPServerTable, + *, + credentials_are_encrypted: bool = True, + env_vars_are_encrypted: bool | None = None, + register_oauth_discovery: bool = True, + ) -> MCPServer: + if mcp_server.server_id == "broken": + raise ValueError("credentials cannot be decrypted") + return await super().build_mcp_server_from_table( + mcp_server, + credentials_are_encrypted=credentials_are_encrypted, + env_vars_are_encrypted=env_vars_are_encrypted, + register_oauth_discovery=register_oauth_discovery, + ) + + +@pytest.mark.asyncio +async def test_one_unbuildable_db_row_does_not_cost_the_other_servers_their_exchange(caplog: pytest.LogCaptureFixture): + catalog = _OneBadRowCatalog({}) + + with caplog.at_level(logging.WARNING, logger="litellm.proxy"): + servers = await login_exchange_servers( + catalog, [_row("broken", MCPAuth.oauth2_id_jag), _row("healthy", MCPAuth.oauth2_id_jag)] + ) + + assert [server.server_id for server in servers] == ["healthy"] + assert any("server_id=broken" in r.getMessage() for r in caplog.records) diff --git a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py index 283d0754940..527475b86a9 100644 --- a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py +++ b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py @@ -32,6 +32,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials import ( Error, IdJagConfig, IdJagResourceClient, + IdJagSubjectSource, NoneConfig, NoOpAuth, Ok, @@ -55,8 +56,11 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_r SSOClientConfig, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( + SAML2_SUBJECT_TOKEN_TYPE, AssertionStoreUnavailable, SSOIdentityAssertion, + assertion_from_saml_login, + with_login_refresh_token, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.token_endpoint import ( ExchangedToken, @@ -111,6 +115,10 @@ class _FakeAssertionStore: self._assertions = dict(assertions or {}) self.lookups: list[str] = [] + async def fetch_stored(self, user_id: str) -> SSOIdentityAssertion | None: + + return await self.fetch(user_id) + async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: self.lookups.append(user_id) return self._assertions.get(user_id) @@ -539,12 +547,8 @@ async def test_id_jag_runs_both_legs_and_returns_the_leg2_bearer(): assert leg1_client_id == "litellm" assert isinstance(endpoint.auths[0], ClientSecretAuth) assert endpoint.auths[0].client_secret.get_secret_value() == "s" - assert ( - leg1_params["grant_type"] == "urn:ietf:params:oauth:grant-type:token-exchange" - ) - assert ( - leg1_params["requested_token_type"] == "urn:ietf:params:oauth:token-type:id-jag" - ) + assert leg1_params["grant_type"] == "urn:ietf:params:oauth:grant-type:token-exchange" + assert leg1_params["requested_token_type"] == "urn:ietf:params:oauth:token-type:id-jag" assert leg1_params["subject_token"] == "user-id-token" leg2_endpoint, leg2_client_id, leg2_params = endpoint.calls[1] @@ -563,9 +567,7 @@ async def test_id_jag_without_inbound_token_or_stored_assertion_is_precondition_ endpoint = _FakeTokenEndpoint([]) store = _FakeAssertionStore() provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) - result = await provider.resolve_credentials( - Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) - ) + result = await provider.resolve_credentials(Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config())) assert isinstance(result, Error) assert result.error.tag == "precondition_required" @@ -581,9 +583,7 @@ async def test_id_jag_exchanges_the_stored_sso_assertion_when_the_caller_present store = _FakeAssertionStore({"alice": _assertion("alice-id-token")}) provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) - result = await provider.resolve_credentials( - Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) - ) + result = await provider.resolve_credentials(Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config())) assert isinstance(result, Ok) assert _emitted(result.ok)["Authorization"] == "Bearer final-access" @@ -593,6 +593,54 @@ async def test_id_jag_exchanges_the_stored_sso_assertion_when_the_caller_present assert leg1_params["requested_token_type"] == "urn:ietf:params:oauth:token-type:id-jag" +@pytest.mark.asyncio +async def test_id_jag_declares_a_stored_saml_assertion_as_a_saml2_subject_token(): + """A SAML login retains the signed assertion, not an id_token, so leg 1 has to tell the IdP what + it is being handed even though the server is configured with the id_token default.""" + endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access")) + saml = assertion_from_saml_login(b"", int(time.time()) + 300, "https://idp.example.com") + store = _FakeAssertionStore({"alice": saml}) + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) + + result = await provider.resolve_credentials(Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config())) + + assert isinstance(result, Ok) + _, _, leg1_params = endpoint.calls[0] + assert leg1_params["subject_token"] == saml.id_token.get_secret_value() + assert leg1_params["subject_token_type"] == SAML2_SUBJECT_TOKEN_TYPE + assert leg1_params["grant_type"] == "urn:ietf:params:oauth:grant-type:token-exchange" + assert leg1_params["requested_token_type"] == "urn:ietf:params:oauth:token-type:id-jag" + + +@pytest.mark.asyncio +async def test_id_jag_sends_the_configured_type_for_a_stored_oidc_assertion(): + endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access")) + store = _FakeAssertionStore({"alice": _assertion("alice-id-token")}) + config = _id_jag_config().model_copy(update={"subject_token_type": "urn:ietf:params:oauth:token-type:jwt"}) + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) + + result = await provider.resolve_credentials(Subject(tenant_id="", subject_id="alice"), _spec(config)) + + assert isinstance(result, Ok) + _, _, leg1_params = endpoint.calls[0] + assert leg1_params["subject_token_type"] == "urn:ietf:params:oauth:token-type:jwt" + + +@pytest.mark.asyncio +async def test_id_jag_refuses_a_saml_assertion_past_its_not_on_or_after_without_calling_the_idp(): + endpoint = _FakeTokenEndpoint([]) + saml = assertion_from_saml_login(b"", int(time.time()) - 1, "https://idp.example.com") + provider = UpstreamCredentialProvider( + token_endpoint=endpoint, sso_assertion_store=_FakeAssertionStore({"alice": saml}) + ) + + result = await provider.resolve_credentials(Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config())) + + assert isinstance(result, Error) + assert result.error.tag == "precondition_required" + assert endpoint.calls == [] + + @pytest.mark.asyncio async def test_id_jag_prefers_the_callers_own_token_over_the_stored_assertion(): endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access")) @@ -613,9 +661,7 @@ async def test_id_jag_refuses_an_expired_stored_assertion_without_calling_the_id store = _FakeAssertionStore({"alice": _assertion("stale-id-token", expires_in=-timedelta(seconds=1))}) provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) - result = await provider.resolve_credentials( - Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) - ) + result = await provider.resolve_credentials(Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config())) assert isinstance(result, Error) assert result.error.tag == "precondition_required" @@ -646,6 +692,9 @@ async def test_id_jag_renews_an_expired_stored_assertion_instead_of_challenging( rows[user_id] = assertion class _Inner: + async def fetch_stored(self, user_id: str) -> SSOIdentityAssertion | None: + return await self.fetch(user_id) + async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: return await _read(user_id) @@ -672,9 +721,7 @@ async def test_id_jag_renews_an_expired_stored_assertion_instead_of_challenging( ), ) - result = await provider.resolve_credentials( - Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) - ) + result = await provider.resolve_credentials(Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config())) assert isinstance(result, Ok) _, _, leg1_params = endpoint.calls[0] @@ -686,7 +733,7 @@ def test_the_resolver_defaults_to_the_renewing_assertion_store(): the plain database reader would strand every agent an id_token lifetime after its user's login.""" provider = UpstreamCredentialProvider() - assert isinstance(provider._sso_assertion_store, RefreshingSSOAssertionStore) # noqa: SLF001 # the wiring is the assertion + assert isinstance(provider._sso_assertion_store, RefreshingSSOAssertionStore) # the wiring is the assertion @pytest.mark.asyncio @@ -695,9 +742,7 @@ async def test_id_jag_accepts_a_stored_assertion_that_declares_no_expiry(): store = _FakeAssertionStore({"alice": _assertion("undated-id-token", expires_in=None)}) provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) - result = await provider.resolve_credentials( - Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) - ) + result = await provider.resolve_credentials(Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config())) assert isinstance(result, Ok) _, _, leg1_params = endpoint.calls[0] @@ -730,14 +775,10 @@ async def test_id_jag_keeps_store_sourced_bearers_partitioned_per_user(): Ok(ExchangedToken(access_token="bob-bearer", expires_in=3600)), ] ) - store = _FakeAssertionStore( - {"alice": _assertion("alice-id-token"), "bob": _assertion("bob-id-token")} - ) + store = _FakeAssertionStore({"alice": _assertion("alice-id-token"), "bob": _assertion("bob-id-token")}) provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) - alice = await provider.resolve_credentials( - Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) - ) + alice = await provider.resolve_credentials(Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config())) bob = await provider.resolve_credentials(Subject(tenant_id="", subject_id="bob"), _spec(_id_jag_config())) assert isinstance(alice, Ok) and isinstance(bob, Ok) @@ -754,6 +795,10 @@ class _OutageAssertionStore: def __init__(self) -> None: self.lookups: list[str] = [] + async def fetch_stored(self, user_id: str) -> SSOIdentityAssertion | None: + + return await self.fetch(user_id) + async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: self.lookups.append(user_id) raise AssertionStoreUnavailable(_DRIVER_DETAIL) @@ -766,9 +811,7 @@ async def test_id_jag_maps_an_assertion_store_outage_to_upstream_unavailable(): endpoint = _FakeTokenEndpoint([]) provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=_OutageAssertionStore()) - result = await provider.resolve_credentials( - Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) - ) + result = await provider.resolve_credentials(Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config())) assert isinstance(result, Error) assert result.error.tag == "upstream_unavailable" @@ -784,9 +827,7 @@ async def test_id_jag_store_outage_does_not_leak_driver_detail_to_the_caller(cap ) with caplog.at_level(logging.WARNING): - result = await provider.resolve_credentials( - Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) - ) + result = await provider.resolve_credentials(Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config())) assert isinstance(result, Error) assert _DRIVER_DETAIL not in result.error.summary @@ -803,9 +844,7 @@ async def test_id_jag_invalidation_survives_an_assertion_store_outage(): token_endpoint=_FakeTokenEndpoint([]), sso_assertion_store=_OutageAssertionStore() ) - await provider.invalidate_credentials( - Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) - ) + await provider.invalidate_credentials(Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config())) class _FlakyAssertionStore: @@ -815,6 +854,10 @@ class _FlakyAssertionStore: self._assertion = assertion self.down = False + async def fetch_stored(self, user_id: str) -> SSOIdentityAssertion | None: + + return await self.fetch(user_id) + async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: if self.down: raise AssertionStoreUnavailable("connection refused") @@ -855,6 +898,10 @@ class _SwitchableAssertionStore: def __init__(self, id_token: str) -> None: self.id_token = id_token + async def fetch_stored(self, user_id: str) -> SSOIdentityAssertion | None: + + return await self.fetch(user_id) + async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: return _assertion(self.id_token) @@ -899,6 +946,10 @@ class _SequentialAssertionStore: self.issued: list[str] = [] self._n = 0 + async def fetch_stored(self, user_id: str) -> SSOIdentityAssertion | None: + + return await self.fetch(user_id) + async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: await asyncio.sleep(0) if self.pinned is not None: @@ -996,14 +1047,10 @@ async def test_id_jag_reads_a_naive_stored_expiry_as_utc(): comparison would be a TypeError on the egress path, turning a 412 into a 500.""" endpoint = _FakeTokenEndpoint([]) naive_past = datetime.now(timezone.utc).replace(tzinfo=None) - timedelta(hours=1) - store = _FakeAssertionStore( - {"alice": SSOIdentityAssertion(id_token=SecretStr("stale"), expires_at=naive_past)} - ) + store = _FakeAssertionStore({"alice": SSOIdentityAssertion(id_token=SecretStr("stale"), expires_at=naive_past)}) provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) - result = await provider.resolve_credentials( - Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) - ) + result = await provider.resolve_credentials(Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config())) assert isinstance(result, Error) assert result.error.tag == "precondition_required" @@ -1032,13 +1079,9 @@ async def test_invalidate_evicts_a_store_sourced_id_jag_bearer(): @pytest.mark.asyncio async def test_id_jag_propagates_a_leg1_error_without_calling_leg2(): - endpoint = _FakeTokenEndpoint( - [Error(CredError.of_upstream_unavailable("leg1 down"))] - ) + endpoint = _FakeTokenEndpoint([Error(CredError.of_upstream_unavailable("leg1 down"))]) provider = UpstreamCredentialProvider(token_endpoint=endpoint) - result = await provider.resolve_credentials( - _with_inbound("user-id-token"), _spec(_id_jag_config()) - ) + result = await provider.resolve_credentials(_with_inbound("user-id-token"), _spec(_id_jag_config())) assert isinstance(result, Error) assert result.error.tag == "upstream_unavailable" @@ -1055,9 +1098,7 @@ async def test_id_jag_propagates_a_leg2_error(): ] ) provider = UpstreamCredentialProvider(token_endpoint=endpoint) - result = await provider.resolve_credentials( - _with_inbound("user-id-token"), _spec(_id_jag_config()) - ) + result = await provider.resolve_credentials(_with_inbound("user-id-token"), _spec(_id_jag_config())) assert isinstance(result, Error) assert result.error.tag == "upstream_unavailable" @@ -1065,7 +1106,7 @@ async def test_id_jag_propagates_a_leg2_error(): assert len(endpoint.calls) == 2 -def _two_leg_ok(bearer: str) -> list: +def _two_leg_ok(bearer: str) -> list[Result[ExchangedToken, CredError]]: return [ Ok(ExchangedToken(access_token="the-id-jag", expires_in=300)), Ok(ExchangedToken(access_token=bearer, expires_in=3600)), @@ -1256,8 +1297,18 @@ async def test_passthrough_ignores_the_carrier_and_keeps_the_callers_slot(): (NoneConfig(), _SUBJECT, "no-auth", None), (PassthroughConfig(), _SUBJECT, "no-auth", None), (PassthroughConfig(), _with_inbound("Bearer caller-token"), "oauth2-passthrough", "Bearer caller-token"), - (ApiKeyConfig(key_source=SharedKey(value=SecretStr("static-key"))), _SUBJECT, "static-token", "Bearer static-key"), - (AuthorizationCodeConfig(), Subject(tenant_id="", subject_id="alice"), "stored-user-token", "Bearer stored-alice"), + ( + ApiKeyConfig(key_source=SharedKey(value=SecretStr("static-key"))), + _SUBJECT, + "static-token", + "Bearer static-key", + ), + ( + AuthorizationCodeConfig(), + Subject(tenant_id="", subject_id="alice"), + "stored-user-token", + "Bearer stored-alice", + ), ], ) async def test_resolved_source_matches_the_credential_sent_upstream( @@ -1279,7 +1330,9 @@ async def test_resolved_source_matches_the_credential_sent_upstream( async def test_resolved_source_preserves_missing_user_token_error() -> None: from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import resolve_credentials_with_source - result = await resolve_credentials_with_source(UpstreamCredentialProvider(), _SUBJECT, _spec(AuthorizationCodeConfig())) + result = await resolve_credentials_with_source( + UpstreamCredentialProvider(), _SUBJECT, _spec(AuthorizationCodeConfig()) + ) assert isinstance(result, Error) assert result.error.tag == "unauthorized" @@ -1289,7 +1342,9 @@ async def test_minted_token_sources_match_egress_and_do_not_fetch_twice() -> Non from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import resolve_credentials_with_source source = _FakeM2MSource(Ok(OAuthToken(access_token="m2m-at"))) - m2m = await resolve_credentials_with_source(UpstreamCredentialProvider(client_credentials_source=source), _SUBJECT, _spec(_M2M)) + m2m = await resolve_credentials_with_source( + UpstreamCredentialProvider(client_credentials_source=source), _SUBJECT, _spec(_M2M) + ) assert isinstance(m2m, Ok) headers, _ = await _emitted_async(m2m.ok.auth) assert headers["Authorization"] == "Bearer m2m-at" @@ -1297,7 +1352,9 @@ async def test_minted_token_sources_match_egress_and_do_not_fetch_twice() -> Non assert source.gets == ["s"] exchanger = _FakeExchanger(Ok(OAuthToken(access_token="exchanged-at"))) - exchanged = await resolve_credentials_with_source(UpstreamCredentialProvider(token_exchanger=exchanger), _with_inbound("subject"), _spec(_OBO)) + exchanged = await resolve_credentials_with_source( + UpstreamCredentialProvider(token_exchanger=exchanger), _with_inbound("subject"), _spec(_OBO) + ) assert isinstance(exchanged, Ok) assert _emitted(exchanged.ok.auth)["Authorization"] == "Bearer exchanged-at" assert exchanged.ok.source.value == "token-exchange" @@ -1321,9 +1378,14 @@ async def test_id_jag_source_describes_final_token_after_both_exchanges() -> Non async def test_client_credentials_subscription_identity_preserves_resolved_grant_scopes(): async def resolve(token): source = _FakeM2MSource(Ok(token)) - result = await UpstreamCredentialProvider(client_credentials_source=source).resolve_credentials(_SUBJECT, _spec(_M2M)) + result = await UpstreamCredentialProvider(client_credentials_source=source).resolve_credentials( + _SUBJECT, _spec(_M2M) + ) assert isinstance(result, Ok) - from litellm.proxy._experimental.mcp_server.outbound_credentials.client_credentials import ClientCredentialsBearerAuth + from litellm.proxy._experimental.mcp_server.outbound_credentials.client_credentials import ( + ClientCredentialsBearerAuth, + ) + assert isinstance(result.ok, ClientCredentialsBearerAuth) return result.ok.subscription_identity @@ -1333,3 +1395,175 @@ async def test_client_credentials_subscription_identity_preserves_resolved_grant assert original is not None assert original == renewed assert original != reduced + + +def _login_with_refresh_token(expires_in: int | None) -> SSOIdentityAssertion: + saml = assertion_from_saml_login(b"", int(time.time()) + 300, "https://idp.example.com") + return with_login_refresh_token( + saml, + "rt-from-login", + expires_in, + datetime.now(timezone.utc), + endpoint="https://idp.example.com/token", + client_id="litellm", + ) + + +def _refresh_token_source_config() -> IdJagConfig: + return _id_jag_config().model_copy(update={"subject_source": IdJagSubjectSource.refresh_token}) + + +@pytest.mark.asyncio +async def test_id_jag_presents_the_login_refresh_token_when_the_server_asks_for_it(): + """``id_jag_subject_source: refresh_token`` presents the refresh token the login assertion was + exchanged for, typed as one, never the SAML assertion itself.""" + endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access")) + store = _FakeAssertionStore({"alice": _login_with_refresh_token(3600)}) + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) + + result = await provider.resolve_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(_refresh_token_source_config()) + ) + + assert isinstance(result, Ok) + _, _, leg1_params = endpoint.calls[0] + assert leg1_params["subject_token"] == "rt-from-login" + assert leg1_params["subject_token_type"] == "urn:ietf:params:oauth:token-type:refresh_token" + assert leg1_params["requested_token_type"] == "urn:ietf:params:oauth:token-type:id-jag" + + +@pytest.mark.asyncio +async def test_id_jag_refresh_token_source_never_presents_the_sso_refresh_token(): + """An OIDC login holds a refresh token for the SSO client. Only the token the ID-JAG client + exchanged at login is presented, the SSO one cannot be redeemed there.""" + endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access")) + oidc = SSOIdentityAssertion(id_token=SecretStr("id-token"), refresh_token=SecretStr("sso-rt")) + row = with_login_refresh_token( + oidc, + "rt-from-login", + 3600, + datetime.now(timezone.utc), + endpoint="https://idp.example.com/token", + client_id="litellm", + ) + provider = UpstreamCredentialProvider( + token_endpoint=endpoint, sso_assertion_store=_FakeAssertionStore({"alice": row}) + ) + + result = await provider.resolve_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(_refresh_token_source_config()) + ) + + assert isinstance(result, Ok) + assert endpoint.calls[0][2]["subject_token"] == "rt-from-login" + + +@pytest.mark.asyncio +async def test_id_jag_refresh_token_source_refuses_a_token_minted_for_another_idp_client(): + """A refresh token is bound to the client that received it. One minted for another server's IdP + client is never posted to this server's endpoint, the user is told to sign in again instead.""" + endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access")) + row = with_login_refresh_token( + _login_with_refresh_token(3600), + "rt-for-other", + 3600, + datetime.now(timezone.utc), + endpoint="https://other-idp.example.com/token", + client_id="other-agent", + ) + provider = UpstreamCredentialProvider( + token_endpoint=endpoint, sso_assertion_store=_FakeAssertionStore({"alice": row}) + ) + + result = await provider.resolve_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(_refresh_token_source_config()) + ) + + assert isinstance(result, Error) + assert result.error.tag == "precondition_required" + assert "different IdP client" in str(result.error) + assert endpoint.calls == [] + + +class _RenewalFailsStore(_FakeAssertionStore): + """An SSO renewal that is timing out: ``fetch`` raises the way RefreshingSSOAssertionStore does, + while the row itself is still readable.""" + + async def fetch_stored(self, user_id: str) -> SSOIdentityAssertion | None: + return self._assertions.get(user_id) + + async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: + raise AssertionStoreUnavailable("IdP token endpoint timed out") + + +@pytest.mark.asyncio +async def test_id_jag_refresh_token_source_does_not_wait_on_an_sso_renewal(): + """The exchanged refresh token has its own lifetime. A failing renewal of the unrelated SSO + id_token must not turn into a 503 for a server that never presents that id_token.""" + endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access")) + provider = UpstreamCredentialProvider( + token_endpoint=endpoint, sso_assertion_store=_RenewalFailsStore({"alice": _login_with_refresh_token(3600)}) + ) + + result = await provider.resolve_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(_refresh_token_source_config()) + ) + + assert isinstance(result, Ok) + assert endpoint.calls[0][2]["subject_token"] == "rt-from-login" + + +@pytest.mark.asyncio +async def test_id_jag_refresh_token_source_outlives_the_saml_assertion(): + """The SAML assertion is spent minutes after login; the refresh token it bought is what counts.""" + endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access")) + expired_saml = assertion_from_saml_login(b"", int(time.time()) - 60, "https://idp.example.com") + row = with_login_refresh_token( + expired_saml, + "rt-from-login", + 3600, + datetime.now(timezone.utc), + endpoint="https://idp.example.com/token", + client_id="litellm", + ) + provider = UpstreamCredentialProvider( + token_endpoint=endpoint, sso_assertion_store=_FakeAssertionStore({"alice": row}) + ) + + result = await provider.resolve_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(_refresh_token_source_config()) + ) + + assert isinstance(result, Ok) + assert endpoint.calls[0][2]["subject_token"] == "rt-from-login" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("row", "expected"), + [ + ( + assertion_from_saml_login(b"", int(time.time()) + 300, "https://idp.example.com"), + "none is stored", + ), + (_login_with_refresh_token(-1), "has expired"), + ], + ids=["no refresh token", "expired refresh token"], +) +async def test_id_jag_refresh_token_source_is_a_precondition_without_a_live_refresh_token( + row: SSOIdentityAssertion, expected: str +): + endpoint = _FakeTokenEndpoint([]) + provider = UpstreamCredentialProvider( + token_endpoint=endpoint, sso_assertion_store=_FakeAssertionStore({"alice": row}) + ) + + result = await provider.resolve_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(_refresh_token_source_config()) + ) + + assert isinstance(result, Error) + assert result.error.tag == "precondition_required" + assert expected in result.error.summary + assert "Sign in through LiteLLM SSO" in result.error.summary + assert endpoint.calls == [] diff --git a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_refresher.py b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_refresher.py index 4197303fdd0..53d0d82a7f2 100644 --- a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_refresher.py +++ b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_refresher.py @@ -35,8 +35,10 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_r sso_client_config, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( + with_login_refresh_token, AssertionStoreUnavailable, SSOIdentityAssertion, + assertion_from_saml_login, ) SIGNING_KEY = "test-idp-signing-key-32-bytes-long-xxxx" @@ -318,6 +320,53 @@ async def test_a_missing_refresh_token_names_the_scope_the_operator_has_to_set(c assert "offline_access" in caplog.text +@pytest.mark.asyncio +async def test_an_expired_saml_assertion_is_never_redeemed_and_names_the_re_login(caplog: pytest.LogCaptureFixture): + """SAML has no refresh grant. Even a row that somehow carries a refresh token must not be posted + to the generic OIDC client, and the operator hint about GENERIC_SCOPE does not apply.""" + saml = assertion_from_saml_login(b"", int(time.time()) - 1, "https://idp.example.com") + with_refresh_token = saml.model_copy(update={"refresh_token": SecretStr("rt_saml")}) + rows = _FakeRows({"alice": with_refresh_token}) + transport = _FakeTransport(_renewal(_id_token())) + + served = await _store(rows, transport).fetch("alice") + refresher = SSOAssertionRefresher(transport, client_config=lambda: _CLIENT, read=rows.fetch, write=rows.write) + with caplog.at_level(logging.WARNING): + outcome = await refresher.refresh("alice", with_refresh_token) + + assert served == with_refresh_token + assert transport.calls == [] + assert rows.writes == [] + match outcome: + case Error(failure): + assert failure.kind == "rejected" + assert "sign in" in failure.detail + case Ok(_): + pytest.fail("a SAML assertion must not be renewable") + assert "sign in" in caplog.text + assert "GENERIC_SCOPE" not in caplog.text + + +class _ForbiddenCoordinator: + async def run(self, user_id: str, key: str, *, refresh: object, reread: object) -> None: + raise AssertionError(f"a SAML assertion must never enter the refresh path (user_id={user_id}, key={key})") + + +@pytest.mark.asyncio +async def test_a_saml_assertion_inside_the_expiry_skew_is_served_as_is_without_electing_a_renewal(): + """SAML cannot be renewed, so a still-valid assertion near its NotOnOrAfter is handed out untouched + instead of running the single-flight election that would 503 the losers.""" + saml = assertion_from_saml_login(b"", int(time.time()) + 30, "https://idp.example.com") + rows = _FakeRows({"alice": saml}) + transport = _FakeTransport(_renewal(_id_token())) + + served = await _store(rows, transport, coordinator_factory=_ForbiddenCoordinator).fetch("alice") + + assert served == saml + assert transport.calls == [] + assert rows.writes == [] + + @pytest.mark.asyncio async def test_an_unconfigured_sso_client_never_calls_the_idp(caplog): rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) @@ -853,3 +902,43 @@ async def test_a_successful_grant_is_handed_back_as_the_parsed_body(): assert isinstance(outcome, Ok) assert outcome.ok["id_token"] == "idt" + + +@pytest.mark.asyncio +async def test_a_renewal_keeps_the_refresh_token_exchanged_at_login_in_the_stored_row(): + """Renewal replaces the id_token and the SSO refresh token. The refresh token the ID-JAG client + exchanged at login is unrelated to either and must still be in the row the next MCP call reads.""" + stale, fresh = _id_token(exp_offset=-1), _id_token() + row = with_login_refresh_token( + _stored(stale, expires_in=-1), + "rt-from-login", + 3600, + datetime.now(timezone.utc), + endpoint="https://idp.example.com/token", + client_id="agent", + ) + rows = _FakeRows({"alice": row}) + + served = await _store(rows, _FakeTransport(_renewal(fresh, "rt_2"))).fetch("alice") + + stored = rows.rows["alice"] + assert stored.id_token.get_secret_value() == fresh + assert stored.refresh_token is not None and stored.refresh_token.get_secret_value() == "rt_2" + assert stored.exchanged_refresh_token is not None + assert stored.exchanged_refresh_token.get_secret_value() == "rt-from-login" + assert stored.exchanged_refresh_token_client_id == "agent" + assert served is not None and served.exchanged_refresh_token is not None + + +@pytest.mark.asyncio +async def test_fetch_stored_hands_back_the_row_without_renewing_it(): + """The read a refresh-token server uses: the row as stored, no IdP round trip even when the SSO + id_token is about to expire.""" + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) + transport = _FakeTransport(_renewal(_id_token())) + + served = await _store(rows, transport).fetch_stored("alice") + + assert served is not None + assert served.id_token.get_secret_value() == rows.rows["alice"].id_token.get_secret_value() + assert transport.calls == [] diff --git a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py index 088b232f3b6..af62594e3b8 100644 --- a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py +++ b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py @@ -8,9 +8,11 @@ rotation re-encrypts stored rows like the sibling per-user credential tables. """ import asyncio +import base64 import json import os import time +from datetime import datetime, timedelta, timezone from unittest.mock import AsyncMock, MagicMock, patch import jwt as pyjwt @@ -18,15 +20,20 @@ import pytest from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( _ASSERTION_CACHE, + SAML2_SUBJECT_TOKEN_TYPE, AssertionStoreUnavailable, DbSSOAssertionStore, SSOAssertionCache, + assertion_from_saml_login, assertion_from_sso_login, + carry_exchanged_refresh_token, ema_assertion_retention_enabled, fetch_sso_identity_assertion, + login_refresh_token_expired, persist_sso_identity_assertion, retain_sso_identity_assertion_for_ema, rotate_sso_identity_assertions_master_key, + with_login_refresh_token, ) from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper, encrypt_value_helper from litellm.types.mcp import MCPAuth @@ -111,6 +118,41 @@ def test_assertion_from_sso_login_happy_path(): assert assertion.expires_at.timestamp() == pytest.approx(time.time() + 3600, abs=5) +SAML_ASSERTION_XML = ( + b'' + b"https://idp.example.com/metadata" +) + + +def _saml_subject_bytes(encoded: str) -> bytes: + assert set(encoded) <= set("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_"), encoded[:40] + return base64.urlsafe_b64decode(encoded + "=" * (-len(encoded) % 4)) + + +def test_assertion_from_saml_login_carries_the_base64url_assertion_as_a_saml2_subject_token(): + not_on_or_after = int(time.time()) + 300 + assertion = assertion_from_saml_login(SAML_ASSERTION_XML, not_on_or_after, "https://idp.example.com/metadata") + assert _saml_subject_bytes(assertion.id_token.get_secret_value()) == SAML_ASSERTION_XML + assert assertion.subject_token_type == SAML2_SUBJECT_TOKEN_TYPE + assert assertion.refresh_token is None + assert assertion.issuer == "https://idp.example.com/metadata" + assert assertion.expires_at is not None + assert assertion.expires_at.tzinfo is not None + assert assertion.expires_at.timestamp() == not_on_or_after + + +def test_assertion_from_saml_login_without_not_on_or_after_has_no_expiry(): + assertion = assertion_from_saml_login(SAML_ASSERTION_XML, None, None) + assert assertion.expires_at is None + assert assertion.subject_token_type == SAML2_SUBJECT_TOKEN_TYPE + + +def test_oidc_login_assertion_leaves_the_subject_token_type_to_the_operator(): + assertion = assertion_from_sso_login(_make_id_token(), None) + assert assertion is not None + assert assertion.subject_token_type is None + + def test_assertion_repr_never_leaks_token_material(): token = _make_id_token() assertion = assertion_from_sso_login(token, "rt_secret_value") @@ -226,6 +268,37 @@ async def test_persist_and_fetch_round_trip_encrypted_at_rest(): assert json.loads(decrypted)["id_token"] == token +@pytest.mark.asyncio +async def test_saml_assertion_round_trips_with_its_subject_token_type(): + stored: dict[str, str] = {} + prisma = _make_prisma(stored) + assertion = assertion_from_saml_login(SAML_ASSERTION_XML, int(time.time()) + 300, "https://idp.example.com") + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + await persist_sso_identity_assertion("user-saml", assertion) + fetched = await fetch_sso_identity_assertion("user-saml") + assert fetched is not None + assert fetched.subject_token_type == SAML2_SUBJECT_TOKEN_TYPE + assert _saml_subject_bytes(fetched.id_token.get_secret_value()) == SAML_ASSERTION_XML + assert fetched.expires_at == assertion.expires_at + assert fetched.refresh_token is None + assert SAML_ASSERTION_XML.decode() not in stored["user-saml"] + + +@pytest.mark.asyncio +async def test_oidc_row_written_before_subject_token_types_reads_back_without_one(): + stored: dict[str, str] = {} + prisma = _make_prisma(stored) + token = _make_id_token() + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + await persist_sso_identity_assertion("user-a", assertion_from_sso_login(token, None)) + fetched = await fetch_sso_identity_assertion("user-a") + assert fetched is not None + assert fetched.subject_token_type is None + decrypted = decrypt_value_helper(stored["user-a"], "test", exception_type="debug") + assert decrypted is not None + assert "subject_token_type" not in json.loads(decrypted) + + @pytest.mark.asyncio async def test_persist_overwrites_previous_login(): stored = {} @@ -248,7 +321,9 @@ async def test_fetch_serves_second_read_from_cache_without_db_read(): cache = SSOAssertionCache() token = _make_id_token() assertion = assertion_from_sso_login(token, "rt_1") - with patch("litellm.proxy.proxy_server.prisma_client", prisma): # test-quality-ok: fake DB seam has no injection boundary + with patch( + "litellm.proxy.proxy_server.prisma_client", prisma + ): # test-quality-ok: fake DB seam has no injection boundary await persist_sso_identity_assertion("user-a", assertion, cache=cache) first = await fetch_sso_identity_assertion("user-a", cache=cache) second = await fetch_sso_identity_assertion("user-a", cache=cache) @@ -266,7 +341,9 @@ async def test_persist_busts_cache_so_relogin_is_visible_immediately(): cache = SSOAssertionCache() first_token = _make_id_token(exp_offset=100) second_token = _make_id_token(exp_offset=7200) - with patch("litellm.proxy.proxy_server.prisma_client", prisma): # test-quality-ok: fake DB seam has no injection boundary + with patch( + "litellm.proxy.proxy_server.prisma_client", prisma + ): # test-quality-ok: fake DB seam has no injection boundary await persist_sso_identity_assertion("user-a", assertion_from_sso_login(first_token, None), cache=cache) first = await fetch_sso_identity_assertion("user-a", cache=cache) await persist_sso_identity_assertion("user-a", assertion_from_sso_login(second_token, "rt_new"), cache=cache) @@ -295,7 +372,9 @@ async def test_fetch_racing_a_relogin_does_not_cache_the_previous_assertion(): return row prisma.db.litellm_ssoidentityassertion.find_unique = _paused_find_unique - with patch("litellm.proxy.proxy_server.prisma_client", prisma): # test-quality-ok: fake DB seam has no injection boundary + with patch( + "litellm.proxy.proxy_server.prisma_client", prisma + ): # test-quality-ok: fake DB seam has no injection boundary await persist_sso_identity_assertion("user-a", assertion_from_sso_login(first_token, None), cache=cache) racing_fetch = asyncio.create_task(fetch_sso_identity_assertion("user-a", cache=cache)) await db_read_started.wait() @@ -324,7 +403,9 @@ async def test_fetch_missing_row_returns_none(): async def test_fetch_does_not_cache_a_missing_row(): prisma = _make_prisma({}) cache = SSOAssertionCache() - with patch("litellm.proxy.proxy_server.prisma_client", prisma): # test-quality-ok: fake DB seam has no injection boundary + with patch( + "litellm.proxy.proxy_server.prisma_client", prisma + ): # test-quality-ok: fake DB seam has no injection boundary first = await fetch_sso_identity_assertion("nobody", cache=cache) second = await fetch_sso_identity_assertion("nobody", cache=cache) assert first is None @@ -353,7 +434,9 @@ async def test_cached_assertion_expires_after_ttl(): cache = SSOAssertionCache(ttl_seconds=1) first_token = _make_id_token(exp_offset=100) second_token = _make_id_token(exp_offset=7200) - with patch("litellm.proxy.proxy_server.prisma_client", prisma): # test-quality-ok: fake DB seam has no injection boundary + with patch( + "litellm.proxy.proxy_server.prisma_client", prisma + ): # test-quality-ok: fake DB seam has no injection boundary await persist_sso_identity_assertion("user-a", assertion_from_sso_login(first_token, None), cache=cache) first = await fetch_sso_identity_assertion("user-a", cache=cache) await persist_sso_identity_assertion( @@ -477,7 +560,9 @@ async def test_db_store_uses_injected_cache(): prisma = _make_prisma(stored) cache = SSOAssertionCache() token = _make_id_token() - with patch("litellm.proxy.proxy_server.prisma_client", prisma): # test-quality-ok: fake DB seam has no injection boundary + with patch( + "litellm.proxy.proxy_server.prisma_client", prisma + ): # test-quality-ok: fake DB seam has no injection boundary await persist_sso_identity_assertion("alice", assertion_from_sso_login(token, None), cache=cache) store = DbSSOAssertionStore(cache=cache) first = await store.fetch("alice") @@ -496,3 +581,58 @@ async def test_db_store_returns_none_for_a_user_with_no_stored_assertion(): with patch.dict(os.environ, {"LITELLM_SALT_KEY": SALT_KEY}): with patch("litellm.proxy.proxy_server.prisma_client", _make_prisma({})): assert await DbSSOAssertionStore().fetch("nobody") is None + + +@pytest.mark.asyncio +async def test_login_refresh_token_round_trips_beside_the_saml_assertion(): + """A SAML login exchanged for a refresh token keeps the assertion, its type and its expiry, and adds + the refresh token with the lifetime the IdP gave it, all encrypted at rest.""" + stored: dict[str, str] = {} + prisma = _make_prisma(stored) + now = datetime.now(timezone.utc).replace(microsecond=0) + saml = assertion_from_saml_login(SAML_ASSERTION_XML, int(time.time()) + 300, "https://idp.example.com") + assertion = with_login_refresh_token( + saml, "rt-from-login", 3600, now, endpoint="https://idp.example.com/token", client_id="agent" + ) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + await persist_sso_identity_assertion("user-rt", assertion) + fetched = await fetch_sso_identity_assertion("user-rt") + assert fetched is not None + assert fetched.exchanged_refresh_token is not None + assert fetched.exchanged_refresh_token.get_secret_value() == "rt-from-login" + assert fetched.exchanged_refresh_token_expires_at == now + timedelta(seconds=3600) + assert fetched.exchanged_refresh_token_endpoint == "https://idp.example.com/token" + assert fetched.exchanged_refresh_token_client_id == "agent" + assert fetched.refresh_token is None + assert fetched.id_token == saml.id_token + assert fetched.expires_at == saml.expires_at + assert fetched.subject_token_type == SAML2_SUBJECT_TOKEN_TYPE + assert "rt-from-login" not in stored["user-rt"] + assert not login_refresh_token_expired(fetched, now) + assert login_refresh_token_expired(fetched, now + timedelta(seconds=3600)) + + +def test_renewal_carries_the_exchanged_refresh_token_forward(): + """An OIDC renewal rebuilds the row from the new id_token. The refresh token the ID-JAG client + exchanged at login is unrelated to that renewal and must survive it.""" + now = datetime.now(timezone.utc).replace(microsecond=0) + previous = with_login_refresh_token( + assertion_from_sso_login(_make_id_token(), "sso-rt"), + "rt-from-login", + 3600, + now, + endpoint="https://idp.example.com/token", + client_id="agent", + ) + renewed = assertion_from_sso_login(_make_id_token(7200), "sso-rt-2") + assert previous is not None and renewed is not None + + carried = carry_exchanged_refresh_token(previous, renewed) + + assert carried.id_token == renewed.id_token + assert carried.refresh_token is not None and carried.refresh_token.get_secret_value() == "sso-rt-2" + assert carried.exchanged_refresh_token is not None + assert carried.exchanged_refresh_token.get_secret_value() == "rt-from-login" + assert carried.exchanged_refresh_token_expires_at == now + timedelta(seconds=3600) + assert carried.exchanged_refresh_token_endpoint == "https://idp.example.com/token" + assert carried.exchanged_refresh_token_client_id == "agent" diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 4961ca0cf42..40901d1dd96 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -1641,7 +1641,8 @@ class TestMCPServerManager: assert server.oauth2_flow is None @pytest.mark.asyncio - async def test_load_servers_from_config_warns_for_id_jag_with_saml_sso(self, monkeypatch, caplog): + async def test_load_servers_from_config_is_silent_for_id_jag_with_saml_sso(self, monkeypatch, caplog): + """A SAML login retains the signed assertion as the saml2 subject token, so SAML is not a gap.""" self._clear_sso_env(monkeypatch) monkeypatch.setenv("SAML_IDP_METADATA_URL", "https://idp.example.com/metadata") manager = MCPServerManager() @@ -1652,11 +1653,7 @@ class TestMCPServerManager: ): await manager.load_servers_from_config(self._id_jag_config()) - warnings = [message for message in caplog.messages if "oauth2_id_jag" in message] - assert len(warnings) == 1 - assert "idjag_server" in warnings[0] - assert "saml" in warnings[0] - assert "MICROSOFT_CLIENT_ID" in warnings[0] + assert [message for message in caplog.messages if "oauth2_id_jag" in message] == [] @pytest.mark.asyncio @pytest.mark.parametrize("client_id_env", ["GOOGLE_CLIENT_ID", "MICROSOFT_CLIENT_ID", "GENERIC_CLIENT_ID"]) diff --git a/tests/unit/proxy/management_endpoints/test_id_jag_assertion_capture.py b/tests/unit/proxy/management_endpoints/test_id_jag_assertion_capture.py index 2218473de6b..97c36fd5792 100644 --- a/tests/unit/proxy/management_endpoints/test_id_jag_assertion_capture.py +++ b/tests/unit/proxy/management_endpoints/test_id_jag_assertion_capture.py @@ -74,14 +74,11 @@ class TestIdJagAssertionCaptureGap: monkeypatch.setenv(env_var, "configured") assert id_jag_assertion_capture_gap() is None - def test_saml_is_named_with_the_oidc_remedy(self, monkeypatch): + def test_saml_has_no_gap(self, monkeypatch: pytest.MonkeyPatch): + """The SAML callback retains the signed assertion as the saml2 subject token, so a SAML + deployment is a capturing one and must not be told to switch to an OIDC provider.""" monkeypatch.setenv("SAML_IDP_METADATA_URL", "https://idp.example.com/metadata") - gap = id_jag_assertion_capture_gap() - assert gap is not None - assert "saml" in gap - assert "GOOGLE_CLIENT_ID" in gap - assert "MICROSOFT_CLIENT_ID" in gap - assert "GENERIC_CLIENT_ID" in gap + assert id_jag_assertion_capture_gap() is None def test_no_sso_configured_reports_a_gap(self): gap = id_jag_assertion_capture_gap() @@ -100,12 +97,9 @@ class TestIdJagAssertionCaptureGapAtStartup: def test_no_provider_at_startup_is_not_yet_a_gap(self): assert id_jag_assertion_capture_gap_at_startup() is None - def test_saml_at_startup_reports_the_capture_gap(self, monkeypatch): + def test_saml_at_startup_has_no_gap(self, monkeypatch: pytest.MonkeyPatch): monkeypatch.setenv("SAML_IDP_METADATA_URL", "https://idp.example.com/metadata") - startup_gap = id_jag_assertion_capture_gap_at_startup() - callback_gap = id_jag_assertion_capture_gap() - assert startup_gap is not None - assert startup_gap == callback_gap + assert id_jag_assertion_capture_gap_at_startup() is None @pytest.mark.parametrize("env_var", ["GOOGLE_CLIENT_ID", "MICROSOFT_CLIENT_ID", "GENERIC_CLIENT_ID"]) def test_an_oidc_provider_at_startup_has_no_gap(self, monkeypatch, env_var): diff --git a/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py index cd5d5e475dd..0d59c655835 100644 --- a/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -4530,7 +4530,6 @@ class TestIdJagRegistrationWarnsAboutTheSSOGap: @pytest.mark.parametrize( "provider_env, expected_fragment", [ - ({"SAML_IDP_METADATA_URL": "https://idp.example.com/metadata"}, "saml"), ({}, "no SSO provider is configured"), ], ) @@ -4564,7 +4563,8 @@ class TestIdJagRegistrationWarnsAboutTheSSOGap: ) self._clear_sso_env(monkeypatch) - monkeypatch.setenv("SAML_IDP_METADATA_URL", "https://idp.example.com/metadata") + for var in ("MICROSOFT_CLIENT_ID", "GOOGLE_CLIENT_ID", "GENERIC_CLIENT_ID", "SAML_IDP_METADATA_URL"): + monkeypatch.delenv(var, raising=False) mock_manager = MagicMock() mock_manager.update_server = AsyncMock() @@ -4601,7 +4601,7 @@ class TestIdJagRegistrationWarnsAboutTheSSOGap: warnings = self._id_jag_warnings(caplog) assert len(warnings) == 1 - assert "saml" in str(warnings[0]) + assert "no SSO provider is configured" in str(warnings[0]) class TestHealthCheckServers: diff --git a/tests/unit/proxy/management_endpoints/test_saml_sso.py b/tests/unit/proxy/management_endpoints/test_saml_sso.py index 01f8fc929a3..181c5e65dae 100644 --- a/tests/unit/proxy/management_endpoints/test_saml_sso.py +++ b/tests/unit/proxy/management_endpoints/test_saml_sso.py @@ -9,15 +9,17 @@ makes a test fail. import base64 import datetime +import logging +import re import time from http.cookies import SimpleCookie import pytest from fastapi import HTTPException, Request -pytest.importorskip( - "onelogin", reason="python3-saml (saml extra) is required for SAML SSO tests" -) +pytest.importorskip("onelogin", reason="python3-saml (saml extra) is required for SAML SSO tests") + +from typing import cast from cryptography import x509 from cryptography.hazmat.primitives import hashes, serialization @@ -26,11 +28,10 @@ from cryptography.x509.oid import NameOID from onelogin.saml2.utils import OneLogin_Saml2_Utils from starlette.datastructures import URL -from typing import cast - from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache +from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import SSOIdentityAssertion from litellm.proxy._types import LitellmUserRoles from litellm.proxy.management_endpoints.sso.saml_sso import ( _SAML_AUTHN_REQUEST_CACHE_PREFIX, @@ -40,15 +41,17 @@ from litellm.proxy.management_endpoints.sso.saml_sso import ( _SAML_REPLAY_GUARD_MAX_TTL_SECONDS, SAMLAuthHandler, ) +from litellm.proxy.management_endpoints.types import CustomOpenID -def _shared_cache(store=None): +def _shared_cache(store: InMemoryCache | None = None) -> DualCache: """A DualCache whose replay guard is backed by a shared, atomic store. An InMemoryCache instance stands in for Redis; passing the same instance to two DualCaches simulates two workers sharing one atomic backend.""" return DualCache(redis_cache=cast(RedisCache, store or InMemoryCache())) + IDP_ENTITY = "https://idp.example.com/metadata" SP_ENTITY = "https://proxy.example.com/sso/saml/metadata" ACS = "https://proxy.example.com/sso/saml/callback" @@ -56,7 +59,7 @@ SSO_URL = "https://idp.example.com/sso" PROXY_BASE_URL = "https://proxy.example.com" -def _make_idp_keypair(): +def _make_idp_keypair() -> tuple[str, str]: key = rsa.generate_private_key(public_exponent=65537, key_size=2048) name = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "idp.example.com")]) cert = ( @@ -79,9 +82,7 @@ def _make_idp_keypair(): def _idp_metadata_xml(cert_pem): - cert_body = "".join( - line for line in cert_pem.splitlines() if "CERTIFICATE" not in line - ) + cert_body = "".join(line for line in cert_pem.splitlines() if "CERTIFICATE" not in line) return ( '' f'' @@ -101,17 +102,17 @@ def _saml_time(delta_seconds): def _build_signed_response( - key_pem, - cert_pem, + key_pem: str, + cert_pem: str, *, - in_response_to=None, - response_level_in_response_to=True, - email="alice@example.com", - attributes=None, - not_before_delta=-60, - not_on_or_after_delta=300, - sign=True, -): + in_response_to: str | None = None, + response_level_in_response_to: bool = True, + email: str = "alice@example.com", + attributes: dict[str, list[str]] | None = None, + not_before_delta: int = -60, + not_on_or_after_delta: int = 300, + sign: bool = True, +) -> str: if attributes is None: attributes = { "email": [email], @@ -157,9 +158,7 @@ def _build_signed_response( if sign: signed = OneLogin_Saml2_Utils.add_sign(assertion, key_pem, cert_pem) - assertion = (signed.decode() if isinstance(signed, bytes) else signed).replace( - '', "" - ) + assertion = (signed.decode() if isinstance(signed, bytes) else signed).replace('', "") return ( '' @@ -174,11 +173,11 @@ def _build_signed_response( ) -def _b64(xml): +def _b64(xml: str) -> str: return base64.b64encode(xml.encode()).decode() -def _fake_request(cookies=None): +def _fake_request(cookies: dict[str, str] | None = None): return type( "Req", (), @@ -190,14 +189,18 @@ def _fake_request(cookies=None): )() +async def _acs_login( + b64: str, cache: DualCache, cookies: dict[str, str] | None = None +) -> tuple[CustomOpenID, SSOIdentityAssertion | None]: + return await SAMLAuthHandler.handle_acs(_fake_request(cookies), cache, {"SAMLResponse": b64}) + + async def _acs(b64, cache, cookies=None): - return await SAMLAuthHandler.handle_acs( - _fake_request(cookies), cache, {"SAMLResponse": b64} - ) + return (await _acs_login(b64, cache, cookies))[0] @pytest.fixture -def saml_env(monkeypatch): +def saml_env(monkeypatch: pytest.MonkeyPatch) -> tuple[str, str]: key_pem, cert_pem = _make_idp_keypair() monkeypatch.setenv("SAML_IDP_METADATA_XML", _idp_metadata_xml(cert_pem)) monkeypatch.setenv("SAML_SP_ENTITY_ID", SP_ENTITY) @@ -214,7 +217,7 @@ def saml_env(monkeypatch): @pytest.fixture -def saml_env_idp_initiated(saml_env, monkeypatch): +def saml_env_idp_initiated(saml_env: tuple[str, str], monkeypatch: pytest.MonkeyPatch) -> tuple[str, str]: monkeypatch.setenv("SAML_ALLOW_UNSOLICITED", "true") return saml_env @@ -234,6 +237,89 @@ async def test_valid_idp_initiated_login_maps_assertion_to_user(saml_env_idp_ini assert result.provider == "saml" +@pytest.mark.asyncio +async def test_valid_login_retains_the_signed_assertion_as_a_saml2_subject_token( + saml_env_idp_initiated: tuple[str, str], +): + """The retained material is the inner , still carrying a verifiable signature, with + the assertion's own NotOnOrAfter as its expiry and no refresh token, since SAML has none.""" + key_pem, cert_pem = saml_env_idp_initiated + resp = _build_signed_response(key_pem, cert_pem, not_on_or_after_delta=300) + + user, assertion = await _acs_login(_b64(resp), _shared_cache()) + + assert assertion is not None + assert assertion.subject_token_type == "urn:ietf:params:oauth:token-type:saml2" + assert assertion.refresh_token is None + assert assertion.issuer == IDP_ENTITY + retained = _retained_bytes(assertion.id_token.get_secret_value()) + assert retained.startswith(b" str: + match = re.search(r']*NotOnOrAfter="([^"]+)"', response_xml) + assert match is not None + return match.group(1) + + +def _signature_valid(xml: bytes, cert_pem: str) -> bool: + return bool(OneLogin_Saml2_Utils.validate_sign(xml, cert_pem, raise_exceptions=True, xpath="//ds:Signature")) + + +def _retained_bytes(encoded: str) -> bytes: + assert set(encoded) <= set("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_"), encoded[:40] + return base64.urlsafe_b64decode(encoded + "=" * (-len(encoded) % 4)) + + +@pytest.mark.asyncio +async def test_retained_assertion_keeps_non_ascii_subject_bytes_verbatim(saml_env_idp_initiated: tuple[str, str]): + key_pem, cert_pem = saml_env_idp_initiated + resp = _build_signed_response(key_pem, cert_pem, email="zo\u00eb@example.com") + + user, assertion = await _acs_login(_b64(resp), _shared_cache()) + + assert assertion is not None + retained = _retained_bytes(assertion.id_token.get_secret_value()) + assert "zo\u00eb@example.com".encode() in retained + assert b"&#" not in retained + assert _signature_valid(retained, cert_pem) + + +class _UnparseableResponseAuth: + def get_last_response_xml(self) -> bytes: + return b" int | None: + return None + + +def test_assertion_retention_failure_does_not_fail_the_validated_login(caplog: pytest.LogCaptureFixture): + """Retention is optional egress material, so a response the retention step cannot parse logs and yields + no assertion instead of turning an IdP-validated login into a 500.""" + with caplog.at_level(logging.WARNING): + retained = SAMLAuthHandler._retained_assertion(_UnparseableResponseAuth()) + + assert retained is None + assert "could not be retained for oauth2_id_jag" in caplog.text + + +@pytest.mark.asyncio +async def test_retained_assertion_is_the_one_the_user_was_mapped_from(saml_env_idp_initiated: tuple[str, str]): + key_pem, cert_pem = saml_env_idp_initiated + resp = _build_signed_response(key_pem, cert_pem, email="bob@example.com") + + user, assertion = await _acs_login(_b64(resp), _shared_cache()) + + assert user.email == "bob@example.com" + assert assertion is not None + assert b"bob@example.com" in _retained_bytes(assertion.id_token.get_secret_value()) + + @pytest.mark.asyncio async def test_tampered_assertion_is_rejected(saml_env): key_pem, cert_pem = saml_env @@ -269,9 +355,7 @@ async def test_signature_from_untrusted_key_is_rejected(saml_env): @pytest.mark.asyncio async def test_expired_assertion_is_rejected(saml_env): key_pem, cert_pem = saml_env - resp = _build_signed_response( - key_pem, cert_pem, not_before_delta=-7200, not_on_or_after_delta=-3600 - ) + resp = _build_signed_response(key_pem, cert_pem, not_before_delta=-7200, not_on_or_after_delta=-3600) with pytest.raises(HTTPException) as exc: await _acs(_b64(resp), DualCache()) @@ -293,9 +377,7 @@ async def test_sp_initiated_known_request_succeeds_once_then_replay_rejected(sam key_pem, cert_pem = saml_env cache = DualCache() request_id = "_authn_req_known" - cache.set_cache( - key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", value="1", ttl=600 - ) + cache.set_cache(key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", value="1", ttl=600) resp = _build_signed_response(key_pem, cert_pem, in_response_to=request_id) cookies = {_SAML_AUTHN_STATE_COOKIE: request_id} @@ -312,9 +394,7 @@ async def test_sp_initiated_response_not_bound_to_browser_is_rejected(saml_env): key_pem, cert_pem = saml_env cache = DualCache() request_id = "_authn_req_known" - cache.set_cache( - key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", value="1", ttl=600 - ) + cache.set_cache(key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", value="1", ttl=600) resp = _build_signed_response(key_pem, cert_pem, in_response_to=request_id) with pytest.raises(HTTPException) as exc: @@ -322,9 +402,7 @@ async def test_sp_initiated_response_not_bound_to_browser_is_rejected(saml_env): assert exc.value.status_code == 401 with pytest.raises(HTTPException) as exc: - await _acs( - _b64(resp), cache, cookies={_SAML_AUTHN_STATE_COOKIE: "_attacker_request"} - ) + await _acs(_b64(resp), cache, cookies={_SAML_AUTHN_STATE_COOKIE: "_attacker_request"}) assert exc.value.status_code == 401 @@ -338,9 +416,7 @@ async def test_subjectconfirmation_only_in_response_to_without_cookie_is_rejecte key_pem, cert_pem = saml_env_idp_initiated cache = DualCache() request_id = "_authn_req_known" - cache.set_cache( - key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", value="1", ttl=600 - ) + cache.set_cache(key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", value="1", ttl=600) resp = _build_signed_response( key_pem, cert_pem, @@ -358,9 +434,7 @@ async def test_subjectconfirmation_only_in_response_to_with_cookie_succeeds(saml key_pem, cert_pem = saml_env cache = DualCache() request_id = "_authn_req_known" - cache.set_cache( - key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", value="1", ttl=600 - ) + cache.set_cache(key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", value="1", ttl=600) resp = _build_signed_response( key_pem, cert_pem, @@ -368,9 +442,7 @@ async def test_subjectconfirmation_only_in_response_to_with_cookie_succeeds(saml response_level_in_response_to=False, ) - result = await _acs( - _b64(resp), cache, cookies={_SAML_AUTHN_STATE_COOKIE: request_id} - ) + result = await _acs(_b64(resp), cache, cookies={_SAML_AUTHN_STATE_COOKIE: request_id}) assert result.email == "alice@example.com" @@ -415,9 +487,7 @@ async def test_assertion_without_id_is_rejected(saml_env_idp_initiated): return None with pytest.raises(HTTPException) as exc: - await SAMLAuthHandler._enforce_response_binding( - _AuthNoAssertionId(), _shared_cache(), None - ) + await SAMLAuthHandler._enforce_response_binding(_AuthNoAssertionId(), _shared_cache(), None) assert exc.value.status_code == 401 assert "ID" in exc.value.detail @@ -450,9 +520,7 @@ async def test_invalid_email_in_assertion_is_rejected_cleanly(saml_env_idp_initi @pytest.mark.asyncio -async def test_email_less_assertion_rejected_when_domain_restriction_configured( - saml_env_idp_initiated, monkeypatch -): +async def test_email_less_assertion_rejected_when_domain_restriction_configured(saml_env_idp_initiated, monkeypatch): key_pem, cert_pem = saml_env_idp_initiated monkeypatch.setenv("ALLOWED_EMAIL_DOMAINS", "example.com") resp = _build_signed_response( @@ -558,11 +626,7 @@ async def test_build_login_redirect_targets_idp_and_caches_request_id(saml_env): location = redirect.headers["location"] assert location.startswith(SSO_URL) assert "SAMLRequest=" in location - cached = [ - k - for k in cache.in_memory_cache.cache_dict - if k.startswith(_SAML_AUTHN_REQUEST_CACHE_PREFIX) - ] + cached = [k for k in cache.in_memory_cache.cache_dict if k.startswith(_SAML_AUTHN_REQUEST_CACHE_PREFIX)] assert len(cached) == 1 request_id = cached[0].split(":", 1)[1] @@ -761,9 +825,7 @@ class TestSAMLAuthnCookieSecureFlag: {"use_x_forwarded_for": True, "mcp_trusted_proxy_ranges": ["10.0.0.0/8"]}, ) cache = DualCache() - request = _fake_request_with_scheme( - "http", headers={"X-Forwarded-Proto": "https"}, client_host="10.0.0.5" - ) + request = _fake_request_with_scheme("http", headers={"X-Forwarded-Proto": "https"}, client_host="10.0.0.5") redirect = await SAMLAuthHandler.build_login_redirect(request, cache) cookie = redirect.headers["set-cookie"] assert "Secure" in cookie @@ -773,9 +835,32 @@ class TestSAMLAuthnCookieSecureFlag: monkeypatch.delenv("PROXY_BASE_URL", raising=False) monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) cache = DualCache() - request = _fake_request_with_scheme( - "http", headers={"X-Forwarded-Proto": "https"}, client_host="203.0.113.5" - ) + request = _fake_request_with_scheme("http", headers={"X-Forwarded-Proto": "https"}, client_host="203.0.113.5") redirect = await SAMLAuthHandler.build_login_redirect(request, cache) cookie = redirect.headers["set-cookie"] assert "Secure" not in cookie + + +def test_an_encrypted_assertion_is_not_retained_and_the_login_log_says_why(caplog: pytest.LogCaptureFixture): + """An IdP that encrypts its assertion still signs the user in, but the ID-JAG arm has nothing to + present. The login log has to name that, otherwise every later 412 looks like a bug re-login cannot fix.""" + + class _Encrypted: + def get_last_response_xml(self) -> bytes: + return ( + b'' + b'' + b"" + ) + + def get_last_assertion_not_on_or_after(self) -> int | None: + return None + + with caplog.at_level(logging.WARNING, logger="litellm.proxy"): + retained = SAMLAuthHandler._retained_assertion(_Encrypted()) # pyright: ignore[reportArgumentType] # stub auth # pyright: ignore[reportPrivateUsage] # the retention seam under test + + assert retained is None + assert any( + "1 saml:EncryptedAssertion" in r.getMessage() and "oauth2_id_jag" in r.getMessage() for r in caplog.records + ) diff --git a/tests/unit/proxy/management_endpoints/test_ui_sso.py b/tests/unit/proxy/management_endpoints/test_ui_sso.py index 91a8d4f7101..765e0506990 100644 --- a/tests/unit/proxy/management_endpoints/test_ui_sso.py +++ b/tests/unit/proxy/management_endpoints/test_ui_sso.py @@ -10,6 +10,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest +from pydantic import SecretStr import respx from fastapi import HTTPException, Request @@ -9004,7 +9005,8 @@ async def test_debug_page_survives_a_store_outage(monkeypatch, caplog): not the page.""" from litellm.proxy.management_endpoints.ui_sso import warn_if_id_jag_capture_gap - monkeypatch.setenv("SAML_IDP_METADATA_URL", "https://idp.example.com/metadata") + for var in ("MICROSOFT_CLIENT_ID", "GOOGLE_CLIENT_ID", "GENERIC_CLIENT_ID", "SAML_IDP_METADATA_URL"): + monkeypatch.delenv(var, raising=False) retention_check = AsyncMock(side_effect=Exception("db down")) with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): assert await warn_if_id_jag_capture_gap(retention_enabled=retention_check) is None @@ -9130,7 +9132,7 @@ async def test_saml_callback_enforces_free_sso_user_limit_after_validation(): async def _fake_handle_acs(**kwargs): call_order.append("validate") - return CustomOpenID( + return (CustomOpenID( id="dana@litellm.ai", email="dana@litellm.ai", first_name=None, @@ -9140,7 +9142,7 @@ async def test_saml_callback_enforces_free_sso_user_limit_after_validation(): provider="saml", team_ids=[], user_role=None, - ) + ), None) async def _fake_count_billable_users(): call_order.append("count") @@ -9662,6 +9664,81 @@ async def test_redirect_from_openid_persists_assertion_under_canonical_user_id() assert response is not None +@pytest.mark.asyncio +async def test_redirect_from_openid_persists_the_login_exchanged_assertion(): + """A server with ``id_jag_subject_source: refresh_token`` needs the refresh token the login + assertion was exchanged for, so the funnel persists what the login-time exchange returns, not + the raw capture.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( + assertion_from_sso_login, + ) + + assertion = assertion_from_sso_login(_ema_id_token(), "rt_1") + assert assertion is not None + mock_request = MagicMock(spec=Request) + mock_request.scope = {} + mock_request.base_url = "http://localhost:4000/" + mock_request.cookies = {} + + retain_mock = AsyncMock() + exchanged = assertion.model_copy(update={"exchanged_refresh_token": SecretStr("rt-from-login")}) + exchange_mock = AsyncMock(return_value=exchanged) + with ( + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch("litellm.proxy.proxy_server.master_key", "sk-master"), + patch("litellm.proxy.proxy_server.general_settings", {}), + patch("litellm.proxy.proxy_server.premium_user", False), + patch("litellm.proxy.proxy_server.user_custom_sso", None), + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), + patch("litellm.proxy.proxy_server.redis_usage_cache", None), + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + patch( + "litellm.proxy.proxy_server.generate_key_helper_fn", + AsyncMock( + return_value={"token": "sk-ui-key", "user_id": "canonical-user-id"} + ), + ), + patch( + "litellm.proxy.management_endpoints.ui_sso.get_user_info_from_db", + AsyncMock(return_value=None), + ), + patch( + "litellm.proxy.management_endpoints.ui_sso.check_and_update_if_proxy_admin_id", + AsyncMock(return_value="internal_user"), + ), + patch( + "litellm.proxy.management_endpoints.ui_sso.retain_sso_identity_assertion_for_ema", + retain_mock, + ), + patch( + "litellm.proxy.management_endpoints.ui_sso.exchange_login_assertion_if_configured", + exchange_mock, + ), + ): + response = await SSOAuthenticationHandler.get_redirect_response_from_openid( + result=CustomOpenID( + id="raw-idp-subject", + email="u@example.com", + first_name="U", + last_name="Ser", + display_name="U Ser", + provider="generic", + team_ids=[], + user_role=None, + ), + request=mock_request, + received_response=None, + generic_client_id="cid", + ui_access_mode=None, + access_token_payload=None, + jwt_handler=None, + sso_assertion=assertion, + ) + exchange_mock.assert_awaited_once_with(assertion) + retain_mock.assert_awaited_once() + assert retain_mock.await_args.kwargs["assertion"] is exchanged + + @pytest.mark.asyncio @pytest.mark.parametrize( "provider_env, sso_class, payload, needs_graph", @@ -9889,7 +9966,6 @@ def _id_jag_gap_warnings(caplog) -> list[str]: @pytest.mark.parametrize( "provider_env, expected_fragment", [ - ({"SAML_IDP_METADATA_URL": "https://idp.example.com/metadata"}, "saml"), ({}, "no SSO provider is configured"), ], ) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 05a4893ad7f..b77399c963f 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -37292,6 +37292,8 @@ export interface components { id_jag_resource_client_secret?: string | null; /** Id Jag Resource Token Endpoint */ id_jag_resource_token_endpoint?: string | null; + /** Id Jag Subject Source */ + id_jag_subject_source?: string | null; /** Redirect Uris */ redirect_uris?: string[] | null; /** Scopes */ From a1c9fa951a5840e0a8395163448e140250c421f3 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Fri, 16 Oct 2026 12:29:33 -0700 Subject: [PATCH 15/58] test(rust-bridge): relocate dispatch tests before routing changes (#46926) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/test-rust.yml | 8 ++++++++ .github/workflows/test-unit.yml | 1 - .../chat_completions/test_dispatch.py | 0 .../embeddings}/__init__.py | 0 .../embeddings/test_dispatch.py | 0 .../{unit => test_litellm_rust}/messages/test_dispatch.py | 0 .../ocr/test_public_dispatch.py} | 0 .../responses/test_dispatch.py | 0 .../rust_bridge => test_litellm_rust}/test_dispatch.py | 0 tests/unit/embeddings/__init__.py | 0 tests/unit/messages/__init__.py | 0 11 files changed, 8 insertions(+), 1 deletion(-) rename tests/{unit => test_litellm_rust}/chat_completions/test_dispatch.py (100%) rename tests/{unit/chat_completions => test_litellm_rust/embeddings}/__init__.py (100%) rename tests/{unit => test_litellm_rust}/embeddings/test_dispatch.py (100%) rename tests/{unit => test_litellm_rust}/messages/test_dispatch.py (100%) rename tests/{unit/ocr/test_dispatch.py => test_litellm_rust/ocr/test_public_dispatch.py} (100%) rename tests/{unit => test_litellm_rust}/responses/test_dispatch.py (100%) rename tests/{unit/rust_bridge => test_litellm_rust}/test_dispatch.py (100%) delete mode 100644 tests/unit/embeddings/__init__.py delete mode 100644 tests/unit/messages/__init__.py diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index f8a7c38ef79..636a5031e5f 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -16,6 +16,10 @@ on: - "litellm/proxy/guardrails/**" - "litellm/utils.py" - "litellm/ocr/**" + - "litellm/chat_completions/dispatch.py" + - "litellm/embeddings/dispatch.py" + - "litellm/messages/dispatch.py" + - "litellm/responses/dispatch.py" - "litellm/llms/base_llm/ocr/**" - "litellm/llms/custom_httpx/llm_http_handler.py" - "Makefile" @@ -46,6 +50,10 @@ on: - "litellm/proxy/guardrails/**" - "litellm/utils.py" - "litellm/ocr/**" + - "litellm/chat_completions/dispatch.py" + - "litellm/embeddings/dispatch.py" + - "litellm/messages/dispatch.py" + - "litellm/responses/dispatch.py" - "litellm/llms/base_llm/ocr/**" - "litellm/llms/custom_httpx/llm_http_handler.py" - "Makefile" diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 0558875bdad..1b486d78f04 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -696,7 +696,6 @@ jobs: tests/unit/responses/litellm_completion_transformation tests/unit/responses/test_additional_tools.py tests/unit/responses/test_custom_tool_call.py - tests/unit/responses/test_dispatch.py tests/unit/responses/test_metadata_codex_callback.py tests/unit/responses/test_no_duplicate_spend_logs.py tests/unit/responses/test_null_test_fix.py diff --git a/tests/unit/chat_completions/test_dispatch.py b/tests/test_litellm_rust/chat_completions/test_dispatch.py similarity index 100% rename from tests/unit/chat_completions/test_dispatch.py rename to tests/test_litellm_rust/chat_completions/test_dispatch.py diff --git a/tests/unit/chat_completions/__init__.py b/tests/test_litellm_rust/embeddings/__init__.py similarity index 100% rename from tests/unit/chat_completions/__init__.py rename to tests/test_litellm_rust/embeddings/__init__.py diff --git a/tests/unit/embeddings/test_dispatch.py b/tests/test_litellm_rust/embeddings/test_dispatch.py similarity index 100% rename from tests/unit/embeddings/test_dispatch.py rename to tests/test_litellm_rust/embeddings/test_dispatch.py diff --git a/tests/unit/messages/test_dispatch.py b/tests/test_litellm_rust/messages/test_dispatch.py similarity index 100% rename from tests/unit/messages/test_dispatch.py rename to tests/test_litellm_rust/messages/test_dispatch.py diff --git a/tests/unit/ocr/test_dispatch.py b/tests/test_litellm_rust/ocr/test_public_dispatch.py similarity index 100% rename from tests/unit/ocr/test_dispatch.py rename to tests/test_litellm_rust/ocr/test_public_dispatch.py diff --git a/tests/unit/responses/test_dispatch.py b/tests/test_litellm_rust/responses/test_dispatch.py similarity index 100% rename from tests/unit/responses/test_dispatch.py rename to tests/test_litellm_rust/responses/test_dispatch.py diff --git a/tests/unit/rust_bridge/test_dispatch.py b/tests/test_litellm_rust/test_dispatch.py similarity index 100% rename from tests/unit/rust_bridge/test_dispatch.py rename to tests/test_litellm_rust/test_dispatch.py diff --git a/tests/unit/embeddings/__init__.py b/tests/unit/embeddings/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/unit/messages/__init__.py b/tests/unit/messages/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 From 1451aed72647caa14cccb72688256f0020ab4324 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Fri, 16 Oct 2026 12:39:26 -0700 Subject: [PATCH 16/58] fix(tests): drain SDK logging before restoring callbacks (#46924) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../database/test_harness_lifecycle.py | 76 +++++++++++++++++++ tests/integration/sdk/conftest.py | 15 +++- 2 files changed, 88 insertions(+), 3 deletions(-) diff --git a/tests/integration/database/test_harness_lifecycle.py b/tests/integration/database/test_harness_lifecycle.py index c8de39c0d38..60bd7ae72a3 100644 --- a/tests/integration/database/test_harness_lifecycle.py +++ b/tests/integration/database/test_harness_lifecycle.py @@ -106,6 +106,82 @@ def test_runner_environment_drops_inherited_provider_api_keys() -> None: assert inherited["OPENAI_API_KEY"] == "ambient-openai" +@pytest.mark.parametrize("fails", (False, True), ids=("passing-test", "failing-test")) +def test_sdk_teardown_delivers_logging_before_restoring_callbacks(tmp_path: Path, fails: bool) -> None: + root: Final = Path(__file__).resolve().parents[3] + events: Final = tmp_path / "callbacks.json" + probe: Final = tmp_path / "test_callback_teardown.py" + probe.write_text( + "import asyncio, json, litellm, pytest\n" + "from pathlib import Path\n" + "from litellm.integrations.custom_logger import CustomLogger\n" + f"EVENTS = Path({str(events)!r})\n" + "ORIGINAL = tuple(litellm.callbacks)\n" + "@pytest.fixture\n" + "def release():\n" + " gate = asyncio.Event()\n" + " yield gate\n" + " gate.set()\n" + "async def test_request(release):\n" + " started = asyncio.Event()\n" + " request_loop = asyncio.get_running_loop()\n" + " class Recorder(CustomLogger):\n" + " async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):\n" + " started.set()\n" + " await release.wait()\n" + " EVENTS.write_text(json.dumps({\n" + " 'registered': self in litellm.callbacks,\n" + " 'same_loop': asyncio.get_running_loop() is request_loop,\n" + " 'call_id': kwargs['litellm_call_id'],\n" + " }))\n" + " litellm.callbacks.append(Recorder())\n" + " response = await litellm.acompletion(model='gpt-4o-mini',\n" + " messages=[{'role': 'user', 'content': 'teardown'}],\n" + " mock_response='done', litellm_call_id='teardown-request')\n" + " await asyncio.wait_for(started.wait(), timeout=5)\n" + " assert response.choices[0].message.content == 'done'\n" + " assert not EVENTS.exists()\n" + f" assert not {fails}, 'intentional test failure'\n" + "def test_next():\n" + " assert EVENTS.exists(), 'logging did not finish before the next test'\n" + " assert json.loads(EVENTS.read_text()) == {\n" + " 'registered': True, 'same_loop': True, 'call_id': 'teardown-request'}\n" + " assert tuple(litellm.callbacks) == ORIGINAL\n" + ) + result: Final = subprocess.run( + ( + sys.executable, + "-P", + "-m", + "pytest", + str(probe), + "--noconftest", + "-p", + "tests.integration.sdk.conftest", + "-o", + "addopts=", + "-o", + "asyncio_mode=auto", + "-o", + "asyncio_default_fixture_loop_scope=session", + f"--junitxml={tmp_path / 'junit.xml'}", + "-q", + ), + cwd=root, + env=_nested_pytest_environment(os.environ), + capture_output=True, + text=True, + timeout=60, + check=False, + ) + assert result.returncode == (1 if fails else 0), result.stdout + result.stderr + cases: Final = ElementTree.parse(tmp_path / "junit.xml").findall(".//testcase") + assert tuple(case.attrib["name"] for case in cases) == ("test_request", "test_next") + assert all(case.find("error") is None for case in cases), result.stdout + result.stderr + assert (cases[0].find("failure") is not None) == fails, result.stdout + result.stderr + assert cases[1].find("failure") is None, result.stdout + result.stderr + + def test_owned_readiness_rejects_a_healthy_listener_from_another_process( monkeypatch: pytest.MonkeyPatch, tmp_path: Path ) -> None: diff --git a/tests/integration/sdk/conftest.py b/tests/integration/sdk/conftest.py index 066cca6250a..05f4f3d7d3a 100644 --- a/tests/integration/sdk/conftest.py +++ b/tests/integration/sdk/conftest.py @@ -1,7 +1,11 @@ +from collections.abc import AsyncIterator from typing import Final -import litellm import pytest +import pytest_asyncio + +import litellm +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER CALLBACK_LISTS: Final = ( "callbacks", @@ -14,7 +18,12 @@ CALLBACK_LISTS: Final = ( ) -@pytest.fixture(autouse=True) -def isolate_litellm_callback_lists(monkeypatch: pytest.MonkeyPatch) -> None: +@pytest_asyncio.fixture(autouse=True, loop_scope="function") +async def isolate_litellm_callback_lists(monkeypatch: pytest.MonkeyPatch) -> AsyncIterator[None]: for name in CALLBACK_LISTS: monkeypatch.setattr(litellm, name, list(getattr(litellm, name))) + try: + yield + finally: + await GLOBAL_LOGGING_WORKER.flush() + await GLOBAL_LOGGING_WORKER.stop() From c1a8575e58efd0bf80a7856312d6478f785d1286 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 16 Oct 2026 12:39:35 -0700 Subject: [PATCH 17/58] test(integration): remove BUG skip for 1 test that passes on main (#46937) test_two_identical_cache_hit_rejects_write_matching_spend_rows passed every run with the skip removed, so the dropped spend row it was guarding against no longer reproduces Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../observability/test_cache_hit_guardrail_metrics.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/integration/observability/test_cache_hit_guardrail_metrics.py b/tests/integration/observability/test_cache_hit_guardrail_metrics.py index 605130acd38..2ca9e0c6410 100644 --- a/tests/integration/observability/test_cache_hit_guardrail_metrics.py +++ b/tests/integration/observability/test_cache_hit_guardrail_metrics.py @@ -577,7 +577,6 @@ def test_two_identical_cache_hit_rejects_increment_populated_series(gateway: Gat def test_two_identical_cache_hit_rejects_write_matching_spend_rows(gateway: Gateway, tmp_path: Path) -> None: """E1b: both cache-hit rejects land a failure spend row.""" - pytest.skip("BUG: roughly one in four back-to-back cache-hit rejects never lands its LiteLLM_SpendLogs row") marker: Final = uuid.uuid4().hex with _rig(gateway, tmp_path, marker) as rig: text: Final = "cache hit control e1b " + marker From 8227f9bf202aaf005861de31584f0aff3dd10334 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Fri, 16 Oct 2026 12:58:32 -0700 Subject: [PATCH 18/58] test(perf): cap every background Postgres statement at 0.9 calls per request (#46957) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../perf/cases/mixed_size_requests_2rps.py | 2 + tests/e2e/perf/cases/small_bodies_5rps.py | 1 + tests/e2e/perf/setup/datastore.py | 5 +++ tests/e2e/perf/setup/observe.py | 35 ++++++++++++---- tests/e2e/perf/setup/performance_test_case.py | 4 ++ tests/e2e_harness/perf/setup/test_observe.py | 42 +++++++++++++++++++ 6 files changed, 81 insertions(+), 8 deletions(-) diff --git a/tests/e2e/perf/cases/mixed_size_requests_2rps.py b/tests/e2e/perf/cases/mixed_size_requests_2rps.py index f95fb0845c4..89263977fb0 100644 --- a/tests/e2e/perf/cases/mixed_size_requests_2rps.py +++ b/tests/e2e/perf/cases/mixed_size_requests_2rps.py @@ -70,6 +70,8 @@ CASE: Final = PerformanceTestCase( Limit("cpu_seconds_per_request", max=1.26), Limit("postgres_calls_per_request", max=3.5), Limit("postgres_max_statement_per_request", max=0.5), + Limit("postgres_max_background_statement_per_request", max=0.9), Limit("redis_calls_per_request", max=21.0), ), + background_cap_exempt_tables=frozenset({"LiteLLM_DailyGuardrailMetrics"}), ) diff --git a/tests/e2e/perf/cases/small_bodies_5rps.py b/tests/e2e/perf/cases/small_bodies_5rps.py index 9d3484de05d..73772465492 100644 --- a/tests/e2e/perf/cases/small_bodies_5rps.py +++ b/tests/e2e/perf/cases/small_bodies_5rps.py @@ -69,6 +69,7 @@ CASE: Final = PerformanceTestCase( Limit("memory_growth_megabytes", max=100), Limit("postgres_calls_per_request", max=1.8), Limit("postgres_max_statement_per_request", max=0.5), + Limit("postgres_max_background_statement_per_request", max=0.9), Limit("redis_calls_per_request", max=18), Limit("cpu_seconds_per_request", max=0.29), ), diff --git a/tests/e2e/perf/setup/datastore.py b/tests/e2e/perf/setup/datastore.py index b4ad10b5a5a..54c6943a054 100644 --- a/tests/e2e/perf/setup/datastore.py +++ b/tests/e2e/perf/setup/datastore.py @@ -76,6 +76,11 @@ class StatementCalls: return "plumbing" return "awaited" if verb == "SELECT" and tables[0] in _AWAITED_TABLES else "background" + @property + def table(self) -> str: + tables: Final = _TABLE.findall(self.query) + return tables[0] if tables else "" + @property def summary(self) -> str: return re.sub(r"\s+", " ", self.query)[:110] diff --git a/tests/e2e/perf/setup/observe.py b/tests/e2e/perf/setup/observe.py index be64679c9da..f634df7de84 100644 --- a/tests/e2e/perf/setup/observe.py +++ b/tests/e2e/perf/setup/observe.py @@ -135,6 +135,14 @@ class Observation: awaited: Final = tuple(item for item in self.statements or () if item.bucket == "awaited") return max(awaited, key=_item_count, default=None) + def busiest_background_statement(self, exempt_tables: frozenset[str]) -> StatementCalls | None: + """The background statement with the most calls, skipping the case's per-flush exemptions: a rollup + written once per request reads 1.0 here while batched writers stay below it.""" + background: Final = tuple( + item for item in self.statements or () if item.bucket == "background" and item.table not in exempt_tables + ) + return max(background, key=_item_count, default=None) + @property def retried_markers(self) -> frozenset[str]: counts: Final = Counter(seen.marker for seen in self.mock_requests if seen.marker is not None) @@ -202,7 +210,9 @@ def _text_sha1(text: str, *, intervened: bool) -> str: return hashlib.sha1(original.encode()).hexdigest() -def observe(metric: Metric, percentile_rank: int | None, observation: Observation) -> float | None: +def observe( + metric: Metric, percentile_rank: int | None, observation: Observation, exempt_tables: frozenset[str] = frozenset() +) -> float | None: """The observed value of one metric, at a percentile when the metric is per request.""" per_request: dict[str, list[float]] = { "gateway_added_seconds": observation.gateway_added_seconds(), @@ -221,6 +231,11 @@ def observe(metric: Metric, percentile_rank: int | None, observation: Observatio return None busiest: Final = observation.busiest_awaited_statement return observation.count_per_request(0 if busiest is None else busiest.calls) + if metric == "postgres_max_background_statement_per_request": + if observation.statements is None: + return None + busiest_background: Final = observation.busiest_background_statement(exempt_tables) + return observation.count_per_request(0 if busiest_background is None else busiest_background.calls) if metric == "redis_calls_per_request": return observation.count_per_request(calls_in("awaited", observation.redis_commands)) if metric == "redis_background_calls_per_request": @@ -266,6 +281,7 @@ REPORTED: Final[tuple[tuple[Metric, int | None], ...]] = ( ("postgres_calls_per_request", None), ("postgres_background_calls_per_request", None), ("postgres_max_statement_per_request", None), + ("postgres_max_background_statement_per_request", None), ("redis_calls_per_request", None), ("redis_background_calls_per_request", None), ("cpu_seconds_per_request", None), @@ -287,9 +303,9 @@ def evaluate(case: PerformanceTestCase, observation: Observation) -> tuple[list[ Row( metric=metric, percentile=rank, - observed=observe(metric, rank, observation), + observed=observe(metric, rank, observation, case.background_cap_exempt_tables), limit=limit.max if limit else None, - detail=_detail(metric, observation), + detail=_detail(metric, observation, case.background_cap_exempt_tables), ) ) exceeded: Final = [ @@ -302,11 +318,14 @@ def evaluate(case: PerformanceTestCase, observation: Observation) -> tuple[list[ return rows, exceeded -def _detail(metric: Metric, observation: Observation) -> str: - if metric != "postgres_max_statement_per_request": - return "" - busiest: Final = observation.busiest_awaited_statement - return "" if busiest is None else busiest.summary +def _detail(metric: Metric, observation: Observation, exempt_tables: frozenset[str]) -> str: + if metric == "postgres_max_statement_per_request": + busiest: Final = observation.busiest_awaited_statement + return "" if busiest is None else busiest.summary + if metric == "postgres_max_background_statement_per_request": + busiest_background: Final = observation.busiest_background_statement(exempt_tables) + return "" if busiest_background is None else busiest_background.summary + return "" def _row_properties(row: Row) -> tuple[tuple[str, float], ...]: diff --git a/tests/e2e/perf/setup/performance_test_case.py b/tests/e2e/perf/setup/performance_test_case.py index 6a734e2923f..37ad640db62 100644 --- a/tests/e2e/perf/setup/performance_test_case.py +++ b/tests/e2e/perf/setup/performance_test_case.py @@ -24,6 +24,7 @@ Metric = Literal[ "postgres_calls_per_request", "postgres_background_calls_per_request", "postgres_max_statement_per_request", + "postgres_max_background_statement_per_request", "cpu_seconds_per_request", "redis_calls_per_request", "redis_background_calls_per_request", @@ -149,6 +150,9 @@ class PerformanceTestCase: traffic: Traffic limits: tuple[Limit, ...] mock_latency: MockLatency = field(default_factory=lambda: MockLatency()) + background_cap_exempt_tables: frozenset[str] = frozenset() + """Tables whose background write is one upsert per writer flush and, at this case's request rate, lands + above the `postgres_max_background_statement_per_request` cap without being per request.""" def team(self, name: str) -> Team: return next(team for team in self.teams if team.name == name) diff --git a/tests/e2e_harness/perf/setup/test_observe.py b/tests/e2e_harness/perf/setup/test_observe.py index 4b7e2d64a14..ad54c61aed0 100644 --- a/tests/e2e_harness/perf/setup/test_observe.py +++ b/tests/e2e_harness/perf/setup/test_observe.py @@ -70,6 +70,7 @@ def test_junit_properties_carry_each_observed_metric_and_each_limit_including_un "perf_postgres_calls_per_request": 10.0, "perf_postgres_background_calls_per_request": 0.0, "perf_postgres_max_statement_per_request": 10.0, + "perf_postgres_max_background_statement_per_request": 0.0, "perf_redis_calls_per_request_limit": 35.0, } @@ -93,6 +94,13 @@ def _observation_with(statements: tuple[StatementCalls, ...]) -> Observation: MAX_STATEMENT_CASE: Final = replace(CASE, limits=(Limit("postgres_max_statement_per_request", max=0.5),)) +BACKGROUND_CAP_CASE: Final = replace( + CASE, + limits=(Limit("postgres_max_background_statement_per_request", max=0.9),), + background_cap_exempt_tables=frozenset({"LiteLLM_DailyGuardrailMetrics"}), +) +GUARDRAIL_UPSERT: Final = 'INSERT INTO "LiteLLM_DailyGuardrailMetrics" ("guardrail_id","date") VALUES ($1,$2)' +MODEL_USAGE_UPSERT: Final = 'INSERT INTO "LiteLLM_DailyModelUsage" ("date","model") VALUES ($1,$2)' TEAM_REFRESH: Final = 'SELECT * FROM "LiteLLM_TeamMembership" WHERE user_id = $1' USER_LOOKUP: Final = 'SELECT * FROM "LiteLLM_UserTable" WHERE user_id = $1' @@ -125,3 +133,37 @@ def test_per_statement_limit_without_db_counters_reports_nothing_observed() -> N _, exceeded = evaluate(MAX_STATEMENT_CASE, observation) assert exceeded == ["postgres_max_statement_per_request: limit set but nothing observed"] + + +def test_per_flush_rollups_pass_the_background_cap_and_the_exempt_table_is_skipped() -> None: + statements: Final = ( + _statement(1, GUARDRAIL_UPSERT, calls=1150), + _statement(2, MODEL_USAGE_UPSERT, calls=810), + _statement(3, USER_LOOKUP, calls=1000), + ) + + rows, exceeded = evaluate(BACKGROUND_CAP_CASE, _observation_with(statements)) + + assert exceeded == [] + assert dict(junit_properties(rows))["perf_postgres_max_background_statement_per_request"] == 0.81 + + +def test_a_rollup_written_on_every_request_fails_the_background_cap_and_names_the_statement() -> None: + statements: Final = (_statement(1, GUARDRAIL_UPSERT, calls=1050), _statement(2, MODEL_USAGE_UPSERT, calls=1000)) + + _, exceeded = evaluate(BACKGROUND_CAP_CASE, _observation_with(statements)) + + assert exceeded == [ + f"postgres_max_background_statement_per_request: observed 1.0000 > limit 0.9000 ({MODEL_USAGE_UPSERT})" + ], exceeded + + +def test_background_cap_without_exemptions_counts_the_guardrail_upsert() -> None: + case: Final = replace(BACKGROUND_CAP_CASE, background_cap_exempt_tables=frozenset()) + statements: Final = (_statement(1, GUARDRAIL_UPSERT, calls=1150), _statement(2, MODEL_USAGE_UPSERT, calls=810)) + + _, exceeded = evaluate(case, _observation_with(statements)) + + assert exceeded == [ + f"postgres_max_background_statement_per_request: observed 1.1500 > limit 0.9000 ({GUARDRAIL_UPSERT})" + ], exceeded From d3ed19e36981f8e9b84594ad33abc72cf338423b Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 20:07:41 +0000 Subject: [PATCH 19/58] fix(rust): forward Anthropic headers to DeepSeek (#46959) Co-authored-by: Mateo --- litellm-rust/crates/gateway-inference/src/messages.rs | 2 +- litellm-rust/crates/gateway-inference/tests/messages.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm-rust/crates/gateway-inference/src/messages.rs b/litellm-rust/crates/gateway-inference/src/messages.rs index c7a8e3e3294..d612a042e03 100644 --- a/litellm-rust/crates/gateway-inference/src/messages.rs +++ b/litellm-rust/crates/gateway-inference/src/messages.rs @@ -24,7 +24,7 @@ const OPENROUTER_API_HEADERS: [&str; 2] = [ "x-openrouter-experimental-metadata", ]; const ANTHROPIC_API_HEADER_PROVIDERS: &str = - "anthropic,bedrock,bedrock_mantle,vertex_ai,openrouter"; + "anthropic,bedrock,bedrock_mantle,vertex_ai,openrouter,deepseek"; pub async fn create( State(gateway): State>, diff --git a/litellm-rust/crates/gateway-inference/tests/messages.rs b/litellm-rust/crates/gateway-inference/tests/messages.rs index b0a9baa2fac..bf11e67e0fc 100644 --- a/litellm-rust/crates/gateway-inference/tests/messages.rs +++ b/litellm-rust/crates/gateway-inference/tests/messages.rs @@ -16,7 +16,7 @@ use wiremock::{ #[rstest] #[case::anthropic("anthropic/test-model", "/v1/messages", true)] #[case::azure("azure_ai/test-model", "/anthropic/v1/messages", false)] -#[case::deepseek("deepseek/test-model", "/anthropic/v1/messages", false)] +#[case::deepseek("deepseek/test-model", "/anthropic/v1/messages", true)] #[tokio::test] async fn messages_reaches_the_provider_and_preserves_json_or_sse( #[case] model: &str, From 33b73207e34892b1ab59ed6c709abf8ae0b89344 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 16 Oct 2026 20:14:53 +0000 Subject: [PATCH 20/58] test(cost): rekey vertex gemini, add bedrock mantle gpt and claude-opus-5-5, drop suite-only price fields (#46962) * test(cost): remove price fields the real gemini and azure rows do not carry Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(cost): key vertex gemini cases on the vertex_ai prefixed rows and add gemini-4-argon Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(cost): add current gemini and vertex usage keys to stored responses Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(cost): add bedrock mantle claude-opus-5-5 cost tracking cases Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(cost): add bedrock mantle gpt cost tracking cases Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../cost_calculation/cases/azure.py | 598 +----------------- .../cost_calculation/cases/bedrock_mantle.py | 454 +++++++++++++ .../cost_calculation/cases/gemini.py | 307 +-------- .../cost_calculation/cases/vertex_ai.py | 441 +++++++------ tests/integration/cost_calculation/catalog.py | 1 + .../integration/cost_calculation/cost_map.py | 203 ++++-- .../cost_calculation/cost_tracking_case.py | 5 + 7 files changed, 859 insertions(+), 1150 deletions(-) diff --git a/tests/integration/cost_calculation/cases/azure.py b/tests/integration/cost_calculation/cases/azure.py index 3df038d1e0e..1360f65b802 100644 --- a/tests/integration/cost_calculation/cases/azure.py +++ b/tests/integration/cost_calculation/cases/azure.py @@ -23,8 +23,18 @@ from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParit from integration.cost_calculation.stream_parity.case import StreamParityTestCase, sse_frames -def _response_for_model(response: JsonResponse, model: str) -> JsonResponse: - return response.model_copy(update={"body": {**response.body, "model": model}}) +def _response_for_model( + response: JsonResponse, model: str, *, service_tier: str | None = None +) -> JsonResponse: + return response.model_copy( + update={ + "body": { + **response.body, + "model": model, + **({"service_tier": service_tier} if service_tier is not None else {}), + } + } + ) def _stream_for_model(response: SseResponse, model: str) -> SseResponse: @@ -3618,173 +3628,7 @@ AZURE_GPT_6_LUNA_STREAM: Final = CostTrackingTestCase( ), ) -AZURE_GPT_6_LUNA_SERVICE_TIER_FLEX: Final = CostTrackingTestCase( - name="azure-gpt-6-luna-service_tier_flex", - covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="azure/gpt-6-luna", - request={ - "model": "$MODEL", - "messages": [ - { - "role": "system", - "content": [ - { - "type": "text", - "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", - } - ], - }, - { - "role": "user", - "content": [ - { - "type": "text", - "text": "895f84528ea4 summarize the attached material in one line and name the city weather", - } - ], - }, - ], - "stream": False, - "service_tier": "flex", - "allowed_openai_params": [], - }, - response=JsonResponse( - content_type="application/json", - body={ - "choices": [ - { - "content_filter_results": _GPT_6_CONTENT_FILTER_RESULTS, - "finish_reason": "stop", - "index": 0, - "logprobs": None, - "message": { - "annotations": [], - "content": "scripted answer 895f84528ea4", - "refusal": None, - "role": "assistant", - }, - } - ], - "created": 1791852029, - "id": "chatcmpl-$REQUEST_ID", - "model": "gpt-6-luna-2026-09-22", - "object": "chat.completion", - "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, - "routing": {"serving_pipereplica": "d20260922011332-060efdec-prefill-r0-dp0-prefill"}, - "service_tier": "flex", - "system_fingerprint": None, - "usage": { - "completion_tokens": 412, - "completion_tokens_details": { - "accepted_prediction_tokens": 0, - "audio_tokens": 0, - "reasoning_tokens": 0, - "rejected_prediction_tokens": 0, - }, - "latency_checkpoint": { - "engine_tbt_ms": 4, - "engine_ttft_ms": 131, - "engine_ttlt_ms": 157, - "pre_inference_ms": 55, - "service_tbt_ms": 5, - "service_ttft_ms": 1034, - "service_ttlt_ms": 1058, - "user_visible_ttft_ms": 979, - }, - "prompt_tokens": 1840, - "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, - "total_tokens": 2252, - }, - }, - ), - expected=ExactExpected( - spend=0.0002535, input_cost=0.0001196, output_cost=0.0001339, prompt_tokens=1840, completion_tokens=412 - ), -) -AZURE_GPT_6_LUNA_SERVICE_TIER_PRIORITY: Final = CostTrackingTestCase( - name="azure-gpt-6-luna-service_tier_priority", - covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="azure/gpt-6-luna", - request={ - "model": "$MODEL", - "messages": [ - { - "role": "system", - "content": [ - { - "type": "text", - "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", - } - ], - }, - { - "role": "user", - "content": [ - { - "type": "text", - "text": "705c026752fc summarize the attached material in one line and name the city weather", - } - ], - }, - ], - "stream": False, - "service_tier": "priority", - "allowed_openai_params": [], - }, - response=JsonResponse( - content_type="application/json", - body={ - "choices": [ - { - "content_filter_results": _GPT_6_CONTENT_FILTER_RESULTS, - "finish_reason": "stop", - "index": 0, - "logprobs": None, - "message": { - "annotations": [], - "content": "scripted answer 705c026752fc", - "refusal": None, - "role": "assistant", - }, - } - ], - "created": 1791852029, - "id": "chatcmpl-$REQUEST_ID", - "model": "gpt-6-luna-2026-09-22", - "object": "chat.completion", - "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, - "routing": {"serving_pipereplica": "d20260922011332-060efdec-prefill-r0-dp0-prefill"}, - "service_tier": "default", - "system_fingerprint": None, - "usage": { - "completion_tokens": 412, - "completion_tokens_details": { - "accepted_prediction_tokens": 0, - "audio_tokens": 0, - "reasoning_tokens": 0, - "rejected_prediction_tokens": 0, - }, - "latency_checkpoint": { - "engine_tbt_ms": 4, - "engine_ttft_ms": 131, - "engine_ttlt_ms": 157, - "pre_inference_ms": 55, - "service_tbt_ms": 5, - "service_ttft_ms": 1034, - "service_ttlt_ms": 1058, - "user_visible_ttft_ms": 979, - }, - "prompt_tokens": 1840, - "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, - "total_tokens": 2252, - }, - }, - ), - expected=ExactExpected( - spend=0.000507, input_cost=0.0002392, output_cost=0.0002678, prompt_tokens=1840, completion_tokens=412 - ), -) AZURE_GPT_6_ASTRA_INPUT_TEXT: Final = CostTrackingTestCase( name="azure-gpt-6-astra-input_text", @@ -3997,173 +3841,7 @@ AZURE_GPT_6_ASTRA_STREAM: Final = CostTrackingTestCase( ), ) -AZURE_GPT_6_ASTRA_SERVICE_TIER_FLEX: Final = CostTrackingTestCase( - name="azure-gpt-6-astra-service_tier_flex", - covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="azure/gpt-6-astra", - request={ - "model": "$MODEL", - "messages": [ - { - "role": "system", - "content": [ - { - "type": "text", - "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", - } - ], - }, - { - "role": "user", - "content": [ - { - "type": "text", - "text": "6ff9027e3f7f summarize the attached material in one line and name the city weather", - } - ], - }, - ], - "stream": False, - "service_tier": "flex", - "allowed_openai_params": [], - }, - response=JsonResponse( - content_type="application/json", - body={ - "choices": [ - { - "content_filter_results": _GPT_6_CONTENT_FILTER_RESULTS, - "finish_reason": "stop", - "index": 0, - "logprobs": None, - "message": { - "annotations": [], - "content": "scripted answer 6ff9027e3f7f", - "refusal": None, - "role": "assistant", - }, - } - ], - "created": 1791852029, - "id": "chatcmpl-$REQUEST_ID", - "model": "gpt-6-astra-2026-09-03", - "object": "chat.completion", - "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, - "routing": {"serving_pipereplica": "d20260922011332-060efdec-prefill-r0-dp0-prefill"}, - "service_tier": "flex", - "system_fingerprint": None, - "usage": { - "completion_tokens": 412, - "completion_tokens_details": { - "accepted_prediction_tokens": 0, - "audio_tokens": 0, - "reasoning_tokens": 0, - "rejected_prediction_tokens": 0, - }, - "latency_checkpoint": { - "engine_tbt_ms": 4, - "engine_ttft_ms": 131, - "engine_ttlt_ms": 157, - "pre_inference_ms": 55, - "service_tbt_ms": 5, - "service_ttft_ms": 1034, - "service_ttlt_ms": 1058, - "user_visible_ttft_ms": 979, - }, - "prompt_tokens": 1840, - "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, - "total_tokens": 2252, - }, - }, - ), - expected=ExactExpected( - spend=0.024375, input_cost=0.0115, output_cost=0.012875, prompt_tokens=1840, completion_tokens=412 - ), -) -AZURE_GPT_6_ASTRA_SERVICE_TIER_PRIORITY: Final = CostTrackingTestCase( - name="azure-gpt-6-astra-service_tier_priority", - covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="azure/gpt-6-astra", - request={ - "model": "$MODEL", - "messages": [ - { - "role": "system", - "content": [ - { - "type": "text", - "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", - } - ], - }, - { - "role": "user", - "content": [ - { - "type": "text", - "text": "9a0c038a2ad8 summarize the attached material in one line and name the city weather", - } - ], - }, - ], - "stream": False, - "service_tier": "priority", - "allowed_openai_params": [], - }, - response=JsonResponse( - content_type="application/json", - body={ - "choices": [ - { - "content_filter_results": _GPT_6_CONTENT_FILTER_RESULTS, - "finish_reason": "stop", - "index": 0, - "logprobs": None, - "message": { - "annotations": [], - "content": "scripted answer 9a0c038a2ad8", - "refusal": None, - "role": "assistant", - }, - } - ], - "created": 1791852029, - "id": "chatcmpl-$REQUEST_ID", - "model": "gpt-6-astra-2026-09-03", - "object": "chat.completion", - "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, - "routing": {"serving_pipereplica": "d20260922011332-060efdec-prefill-r0-dp0-prefill"}, - "service_tier": "default", - "system_fingerprint": None, - "usage": { - "completion_tokens": 412, - "completion_tokens_details": { - "accepted_prediction_tokens": 0, - "audio_tokens": 0, - "reasoning_tokens": 0, - "rejected_prediction_tokens": 0, - }, - "latency_checkpoint": { - "engine_tbt_ms": 4, - "engine_ttft_ms": 131, - "engine_ttlt_ms": 157, - "pre_inference_ms": 55, - "service_tbt_ms": 5, - "service_ttft_ms": 1034, - "service_ttlt_ms": 1058, - "user_visible_ttft_ms": 979, - }, - "prompt_tokens": 1840, - "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, - "total_tokens": 2252, - }, - }, - ), - expected=ExactExpected( - spend=0.04875, input_cost=0.023, output_cost=0.02575, prompt_tokens=1840, completion_tokens=412 - ), -) AZURE_GPT_6_1_SOL_INPUT_TEXT: Final = CostTrackingTestCase( name="azure-gpt-6.1-sol-input_text", @@ -4376,173 +4054,7 @@ AZURE_GPT_6_1_SOL_STREAM: Final = CostTrackingTestCase( ), ) -AZURE_GPT_6_1_SOL_SERVICE_TIER_FLEX: Final = CostTrackingTestCase( - name="azure-gpt-6.1-sol-service_tier_flex", - covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="azure/gpt-6.1-sol", - request={ - "model": "$MODEL", - "messages": [ - { - "role": "system", - "content": [ - { - "type": "text", - "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", - } - ], - }, - { - "role": "user", - "content": [ - { - "type": "text", - "text": "960e0935140f summarize the attached material in one line and name the city weather", - } - ], - }, - ], - "stream": False, - "service_tier": "flex", - "allowed_openai_params": [], - }, - response=JsonResponse( - content_type="application/json", - body={ - "choices": [ - { - "content_filter_results": _GPT_6_CONTENT_FILTER_RESULTS, - "finish_reason": "stop", - "index": 0, - "logprobs": None, - "message": { - "annotations": [], - "content": "scripted answer 960e0935140f", - "refusal": None, - "role": "assistant", - }, - } - ], - "created": 1791852029, - "id": "chatcmpl-$REQUEST_ID", - "model": "gpt-6.1-sol-2026-09-29", - "object": "chat.completion", - "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, - "routing": {"serving_pipereplica": "d20260922011332-060efdec-prefill-r0-dp0-prefill"}, - "service_tier": "flex", - "system_fingerprint": None, - "usage": { - "completion_tokens": 412, - "completion_tokens_details": { - "accepted_prediction_tokens": 0, - "audio_tokens": 0, - "reasoning_tokens": 0, - "rejected_prediction_tokens": 0, - }, - "latency_checkpoint": { - "engine_tbt_ms": 4, - "engine_ttft_ms": 131, - "engine_ttlt_ms": 157, - "pre_inference_ms": 55, - "service_tbt_ms": 5, - "service_ttft_ms": 1034, - "service_ttlt_ms": 1058, - "user_visible_ttft_ms": 979, - }, - "prompt_tokens": 1840, - "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, - "total_tokens": 2252, - }, - }, - ), - expected=ExactExpected( - spend=0.00468, input_cost=0.002208, output_cost=0.002472, prompt_tokens=1840, completion_tokens=412 - ), -) -AZURE_GPT_6_1_SOL_SERVICE_TIER_PRIORITY: Final = CostTrackingTestCase( - name="azure-gpt-6.1-sol-service_tier_priority", - covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="azure/gpt-6.1-sol", - request={ - "model": "$MODEL", - "messages": [ - { - "role": "system", - "content": [ - { - "type": "text", - "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", - } - ], - }, - { - "role": "user", - "content": [ - { - "type": "text", - "text": "6b36e2dc1f33 summarize the attached material in one line and name the city weather", - } - ], - }, - ], - "stream": False, - "service_tier": "priority", - "allowed_openai_params": [], - }, - response=JsonResponse( - content_type="application/json", - body={ - "choices": [ - { - "content_filter_results": _GPT_6_CONTENT_FILTER_RESULTS, - "finish_reason": "stop", - "index": 0, - "logprobs": None, - "message": { - "annotations": [], - "content": "scripted answer 6b36e2dc1f33", - "refusal": None, - "role": "assistant", - }, - } - ], - "created": 1791852029, - "id": "chatcmpl-$REQUEST_ID", - "model": "gpt-6.1-sol-2026-09-29", - "object": "chat.completion", - "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, - "routing": {"serving_pipereplica": "d20260922011332-060efdec-prefill-r0-dp0-prefill"}, - "service_tier": "default", - "system_fingerprint": None, - "usage": { - "completion_tokens": 412, - "completion_tokens_details": { - "accepted_prediction_tokens": 0, - "audio_tokens": 0, - "reasoning_tokens": 0, - "rejected_prediction_tokens": 0, - }, - "latency_checkpoint": { - "engine_tbt_ms": 4, - "engine_ttft_ms": 131, - "engine_ttlt_ms": 157, - "pre_inference_ms": 55, - "service_tbt_ms": 5, - "service_ttft_ms": 1034, - "service_ttlt_ms": 1058, - "user_visible_ttft_ms": 979, - }, - "prompt_tokens": 1840, - "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, - "total_tokens": 2252, - }, - }, - ), - expected=ExactExpected( - spend=0.00936, input_cost=0.004416, output_cost=0.004944, prompt_tokens=1840, completion_tokens=412 - ), -) AZURE_GPT_6_SOL_REASONING: Final = CostTrackingTestCase( name="azure/gpt-6-sol-reasoning", @@ -5995,31 +5507,7 @@ AZURE_GPT_6_LUNA_TIERED_AT_272K_BOUNDARY: Final = with_usage( ), ) -AZURE_GPT_6_LUNA_TIERED_FLEX_ABOVE_272K: Final = with_usage( - AZURE_GPT_6_LUNA_SERVICE_TIER_FLEX, - "azure-gpt-6-luna-tiered_flex_above_272k", - { - "prompt_tokens": 300000, - "total_tokens": 300412, - "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, - }, - ExactExpected( - spend=0.0392678, input_cost=0.039, output_cost=0.0002678, prompt_tokens=300000, completion_tokens=412 - ), -) -AZURE_GPT_6_LUNA_TIERED_PRIORITY_ABOVE_272K: Final = with_usage( - AZURE_GPT_6_LUNA_SERVICE_TIER_PRIORITY, - "azure-gpt-6-luna-tiered_priority_above_272k", - { - "prompt_tokens": 300000, - "total_tokens": 300412, - "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, - }, - ExactExpected( - spend=0.0785356, input_cost=0.078, output_cost=0.0005356, prompt_tokens=300000, completion_tokens=412 - ), -) AZURE_GPT_6_ASTRA_TIERED_INPUT_ABOVE_272K: Final = with_usage( AZURE_GPT_6_ASTRA_INPUT_TEXT, @@ -6054,27 +5542,7 @@ AZURE_GPT_6_ASTRA_TIERED_AT_272K_BOUNDARY: Final = with_usage( ExactExpected(spend=3.42575, input_cost=3.4, output_cost=0.02575, prompt_tokens=272000, completion_tokens=412), ) -AZURE_GPT_6_ASTRA_TIERED_FLEX_ABOVE_272K: Final = with_usage( - AZURE_GPT_6_ASTRA_SERVICE_TIER_FLEX, - "azure-gpt-6-astra-tiered_flex_above_272k", - { - "prompt_tokens": 300000, - "total_tokens": 300412, - "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, - }, - ExactExpected(spend=3.77575, input_cost=3.75, output_cost=0.02575, prompt_tokens=300000, completion_tokens=412), -) -AZURE_GPT_6_ASTRA_TIERED_PRIORITY_ABOVE_272K: Final = with_usage( - AZURE_GPT_6_ASTRA_SERVICE_TIER_PRIORITY, - "azure-gpt-6-astra-tiered_priority_above_272k", - { - "prompt_tokens": 300000, - "total_tokens": 300412, - "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, - }, - ExactExpected(spend=7.5515, input_cost=7.5, output_cost=0.0515, prompt_tokens=300000, completion_tokens=412), -) AZURE_GPT_6_1_SOL_TIERED_INPUT_ABOVE_272K: Final = with_usage( AZURE_GPT_6_1_SOL_INPUT_TEXT, @@ -6109,27 +5577,7 @@ AZURE_GPT_6_1_SOL_TIERED_AT_272K_BOUNDARY: Final = with_usage( ExactExpected(spend=0.657744, input_cost=0.6528, output_cost=0.004944, prompt_tokens=272000, completion_tokens=412), ) -AZURE_GPT_6_1_SOL_TIERED_FLEX_ABOVE_272K: Final = with_usage( - AZURE_GPT_6_1_SOL_SERVICE_TIER_FLEX, - "azure-gpt-6.1-sol-tiered_flex_above_272k", - { - "prompt_tokens": 300000, - "total_tokens": 300412, - "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, - }, - ExactExpected(spend=0.724944, input_cost=0.72, output_cost=0.004944, prompt_tokens=300000, completion_tokens=412), -) -AZURE_GPT_6_1_SOL_TIERED_PRIORITY_ABOVE_272K: Final = with_usage( - AZURE_GPT_6_1_SOL_SERVICE_TIER_PRIORITY, - "azure-gpt-6.1-sol-tiered_priority_above_272k", - { - "prompt_tokens": 300000, - "total_tokens": 300412, - "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, - }, - ExactExpected(spend=1.449888, input_cost=1.44, output_cost=0.009888, prompt_tokens=300000, completion_tokens=412), -) AZURE_GPT_6_1_SOL_CACHE_WRITE_5M: Final = with_response_usage( AZURE_GPT_6_1_SOL_CACHE_READ, @@ -6221,12 +5669,12 @@ AZURE_GPT_5_6_TERRA_STREAM: Final = AZURE_GPT_6_LUNA_STREAM.model_copy( ), } ) -AZURE_GPT_5_6_TERRA_SERVICE_TIER_FLEX: Final = AZURE_GPT_6_LUNA_SERVICE_TIER_FLEX.model_copy( +AZURE_GPT_5_6_TERRA_SERVICE_TIER_FLEX: Final = AZURE_GPT_5_6_SERVICE_TIER_FLEX.model_copy( update={ "name": "azure-gpt-5.6-terra-service_tier_flex", "model": "azure/gpt-5.6-terra", "response": _response_for_model( - cast(JsonResponse, AZURE_GPT_6_LUNA_SERVICE_TIER_FLEX.response), "gpt-5.6-terra-2026-07-09" + cast(JsonResponse, AZURE_GPT_5_6_SERVICE_TIER_FLEX.response), "gpt-5.6-terra-2026-07-09" ), "expected": ExactExpected( spend=0.0047432, @@ -6237,12 +5685,14 @@ AZURE_GPT_5_6_TERRA_SERVICE_TIER_FLEX: Final = AZURE_GPT_6_LUNA_SERVICE_TIER_FLE ), } ) -AZURE_GPT_5_6_TERRA_SERVICE_TIER_PRIORITY: Final = AZURE_GPT_6_LUNA_SERVICE_TIER_PRIORITY.model_copy( +AZURE_GPT_5_6_TERRA_SERVICE_TIER_PRIORITY: Final = AZURE_GPT_5_6_SERVICE_TIER_PRIORITY.model_copy( update={ "name": "azure-gpt-5.6-terra-service_tier_priority", "model": "azure/gpt-5.6-terra", "response": _response_for_model( - cast(JsonResponse, AZURE_GPT_6_LUNA_SERVICE_TIER_PRIORITY.response), "gpt-5.6-terra-2026-07-09" + cast(JsonResponse, AZURE_GPT_5_6_SERVICE_TIER_PRIORITY.response), + "gpt-5.6-terra-2026-07-09", + service_tier="default", ), "expected": ExactExpected( spend=0.0189728, @@ -6376,18 +5826,12 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( AZURE_GPT_6_LUNA_TIERED_INPUT_ABOVE_272K, AZURE_GPT_6_LUNA_TIERED_CACHE_READ_ABOVE_272K, AZURE_GPT_6_LUNA_TIERED_AT_272K_BOUNDARY, - AZURE_GPT_6_LUNA_TIERED_FLEX_ABOVE_272K, - AZURE_GPT_6_LUNA_TIERED_PRIORITY_ABOVE_272K, AZURE_GPT_6_ASTRA_TIERED_INPUT_ABOVE_272K, AZURE_GPT_6_ASTRA_TIERED_CACHE_READ_ABOVE_272K, AZURE_GPT_6_ASTRA_TIERED_AT_272K_BOUNDARY, - AZURE_GPT_6_ASTRA_TIERED_FLEX_ABOVE_272K, - AZURE_GPT_6_ASTRA_TIERED_PRIORITY_ABOVE_272K, AZURE_GPT_6_1_SOL_TIERED_INPUT_ABOVE_272K, AZURE_GPT_6_1_SOL_TIERED_CACHE_READ_ABOVE_272K, AZURE_GPT_6_1_SOL_TIERED_AT_272K_BOUNDARY, - AZURE_GPT_6_1_SOL_TIERED_FLEX_ABOVE_272K, - AZURE_GPT_6_1_SOL_TIERED_PRIORITY_ABOVE_272K, AZURE_GPT_5_4_MINI_INPUT_TEXT, AZURE_GPT_5_4_MINI_CACHE_READ, AZURE_GPT_5_4_MINI_AUDIO_INPUT, @@ -6436,16 +5880,10 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( AZURE_GPT_6_SOL_SERVICE_TIER_PRIORITY, AZURE_GPT_6_LUNA_INPUT_TEXT, AZURE_GPT_6_LUNA_STREAM, - AZURE_GPT_6_LUNA_SERVICE_TIER_FLEX, - AZURE_GPT_6_LUNA_SERVICE_TIER_PRIORITY, AZURE_GPT_6_ASTRA_INPUT_TEXT, AZURE_GPT_6_ASTRA_STREAM, - AZURE_GPT_6_ASTRA_SERVICE_TIER_FLEX, - AZURE_GPT_6_ASTRA_SERVICE_TIER_PRIORITY, AZURE_GPT_6_1_SOL_INPUT_TEXT, AZURE_GPT_6_1_SOL_STREAM, - AZURE_GPT_6_1_SOL_SERVICE_TIER_FLEX, - AZURE_GPT_6_1_SOL_SERVICE_TIER_PRIORITY, AZURE_GPT_6_SOL_REASONING, AZURE_GPT_6_SOL_STREAM_NO_USAGE, AZURE_GPT_6_LUNA_REASONING, diff --git a/tests/integration/cost_calculation/cases/bedrock_mantle.py b/tests/integration/cost_calculation/cases/bedrock_mantle.py index 714c6f2811f..2fa940d66f4 100644 --- a/tests/integration/cost_calculation/cases/bedrock_mantle.py +++ b/tests/integration/cost_calculation/cases/bedrock_mantle.py @@ -2,9 +2,23 @@ from __future__ import annotations +import json +from collections.abc import Mapping from typing import Final from integration.cost_calculation.cases.openai import with_usage +from integration.cost_calculation.cases.anthropic import ( + CLAUDE_OPUS_5_5_CACHE_READ, + CLAUDE_OPUS_5_5_CACHE_WRITE_1H, + CLAUDE_OPUS_5_5_CACHE_WRITE_5M, + CLAUDE_OPUS_5_5_INPUT_TEXT, + CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, + CLAUDE_OPUS_5_5_MESSAGES_STREAM, + CLAUDE_OPUS_5_5_STREAM, + CLAUDE_OPUS_5_5_THINKING, + CLAUDE_OPUS_5_5_TOOL_CALL, +) +from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParityTestCase from integration.cost_calculation.cost_tracking_case import ( CostTrackingTestCase, Deployment, @@ -13,6 +27,7 @@ from integration.cost_calculation.cost_tracking_case import ( SseResponse, ) from integration.cost_calculation.stream_parity.case import StreamParityTestCase, sse_frames +from pydantic import JsonValue BEDROCK_MANTLE_OPENAI_GPT_6_SOL_RESPONSES_JSON: Final = CostTrackingTestCase( @@ -503,6 +518,420 @@ BEDROCK_MANTLE_OPENAI_GPT_6_SOL_TIERED_AT_272K_BOUNDARY: Final = with_usage( ExactExpected(spend=0.602932, input_cost=0.5984, output_cost=0.004532, prompt_tokens=272000, completion_tokens=412), ) +_MANTLE_GPT_RESPONSES_REQUEST: Final = { + "model": "$MODEL", + "input": "cost tracking fixture", + "max_output_tokens": 64, + "stream": False, +} +_MANTLE_GPT_REASONING_REQUEST: Final = { + **_MANTLE_GPT_RESPONSES_REQUEST, + "max_output_tokens": 128, + "reasoning": {"effort": "medium"}, +} +_MANTLE_GPT_STREAM_REQUEST: Final = {**_MANTLE_GPT_RESPONSES_REQUEST, "stream": True} +_MANTLE_GPT_CHAT_REQUEST: Final = { + "model": "$MODEL", + "messages": [{"role": "user", "content": "cost tracking fixture"}], + "max_completion_tokens": 64, +} + + +def _mantle_gpt_usage( + input_tokens: int, output_tokens: int, reasoning_tokens: int, cached_tokens: int = 0 +) -> dict[str, JsonValue]: + return { + "input_tokens": input_tokens, + "input_tokens_details": {"cache_write_tokens": 0, "cached_tokens": cached_tokens}, + "output_tokens": output_tokens, + "output_tokens_details": {"reasoning_tokens": reasoning_tokens}, + "total_tokens": input_tokens + output_tokens, + } + + +def _mantle_gpt_json_case( + base: CostTrackingTestCase, + provider_model: str, + shape: str, + usage: Mapping[str, JsonValue], + expected: ExactExpected, + request: Mapping[str, JsonValue], +) -> CostTrackingTestCase: + case: Final = with_usage(base, f"bedrock_mantle-{provider_model}-{shape}", usage, expected) + if not isinstance(case.response, JsonResponse): + raise TypeError(f"{case.name}: expected a JSON response fixture") + response: Final = case.response.model_copy(update={"body": {**case.response.body, "model": provider_model}}) + return case.model_copy( + update={"model": f"bedrock_mantle/{provider_model}", "request": dict(request), "response": response} + ) + + +def _mantle_gpt_stream_case( + provider_model: str, usage: Mapping[str, JsonValue], expected: ExactExpected +) -> CostTrackingTestCase: + base: Final = BEDROCK_MANTLE_OPENAI_GPT_6_SOL_RESPONSES_STREAM + if not isinstance(base.response, SseResponse): + raise TypeError(f"{base.name}: expected an SSE response fixture") + old_usage: Final = json.dumps(_mantle_gpt_usage(1840, 412, 0)) + new_usage: Final = json.dumps(dict(usage)) + frames: Final = tuple( + frame.replace("openai.gpt-6-sol", provider_model).replace(old_usage, new_usage) + for frame in base.response.frames + ) + response: Final = base.response.model_copy(update={"frames": frames}) + return base.model_copy( + update={ + "name": f"bedrock_mantle-{provider_model}-responses_stream", + "model": f"bedrock_mantle/{provider_model}", + "request": dict(_MANTLE_GPT_STREAM_REQUEST), + "response": response, + "expected": expected, + } + ) + + +def _mantle_gpt_cases( + provider_model: str, + response_reasoning_tokens: int, + response_expected: ExactExpected, + reasoning_usage: Mapping[str, JsonValue], + reasoning_expected: ExactExpected, + cache_reasoning_tokens: int, + cache_expected: ExactExpected, + tiered_input_expected: ExactExpected, + tiered_cache_expected: ExactExpected, + tiered_boundary_expected: ExactExpected, +) -> tuple[CostTrackingTestCase, ...]: + response_usage: Final = _mantle_gpt_usage(9, 64, response_reasoning_tokens) + response_case: Final = _mantle_gpt_json_case( + BEDROCK_MANTLE_OPENAI_GPT_6_SOL_RESPONSES_JSON, + provider_model, + "responses_json", + response_usage, + response_expected, + _MANTLE_GPT_RESPONSES_REQUEST, + ) + stream_case: Final = _mantle_gpt_stream_case(provider_model, response_usage, response_expected) + reasoning_case: Final = _mantle_gpt_json_case( + BEDROCK_MANTLE_OPENAI_GPT_6_SOL_RESPONSES_REASONING, + provider_model, + "responses_reasoning", + reasoning_usage, + reasoning_expected, + _MANTLE_GPT_REASONING_REQUEST, + ) + chat_case: Final = _mantle_gpt_json_case( + BEDROCK_MANTLE_OPENAI_GPT_6_SOL_CHAT_JSON, + provider_model, + "chat_json", + response_usage, + response_expected, + _MANTLE_GPT_CHAT_REQUEST, + ) + cache_case: Final = _mantle_gpt_json_case( + BEDROCK_MANTLE_OPENAI_GPT_6_SOL_RESPONSES_JSON, + provider_model, + "responses_cache_read", + _mantle_gpt_usage(3144, 16, cache_reasoning_tokens, cached_tokens=3142), + cache_expected, + {**_MANTLE_GPT_RESPONSES_REQUEST, "max_output_tokens": 16}, + ) + tiered_input: Final = _mantle_gpt_json_case( + BEDROCK_MANTLE_OPENAI_GPT_6_SOL_RESPONSES_JSON, + provider_model, + "tiered_input_above_272k", + _mantle_gpt_usage(300000, 412, 0), + tiered_input_expected, + {**_MANTLE_GPT_RESPONSES_REQUEST, "max_output_tokens": 412}, + ) + tiered_cache: Final = _mantle_gpt_json_case( + BEDROCK_MANTLE_OPENAI_GPT_6_SOL_RESPONSES_JSON, + provider_model, + "tiered_cache_read_above_272k", + _mantle_gpt_usage(300000, 412, 0, cached_tokens=200000), + tiered_cache_expected, + {**_MANTLE_GPT_RESPONSES_REQUEST, "max_output_tokens": 412}, + ) + tiered_boundary: Final = _mantle_gpt_json_case( + BEDROCK_MANTLE_OPENAI_GPT_6_SOL_RESPONSES_JSON, + provider_model, + "tiered_at_272k_boundary", + _mantle_gpt_usage(272000, 412, 0), + tiered_boundary_expected, + {**_MANTLE_GPT_RESPONSES_REQUEST, "max_output_tokens": 412}, + ) + return (response_case, stream_case, reasoning_case, chat_case, cache_case, tiered_input, tiered_cache, tiered_boundary) + + +def _mantle_gpt_bridge(case: CostTrackingTestCase) -> BridgeCostParityTestCase: + bridged: Final = case.model_copy( + update={ + "name": case.name.removesuffix("responses_json") + "chat_bridged_json", + "endpoint": "/v1/chat/completions", + "request": dict(_MANTLE_GPT_CHAT_REQUEST), + } + ) + return BridgeCostParityTestCase(native=case, bridged=bridged) + + +_MANTLE_GPT_61_SOL_CASES: Final = _mantle_gpt_cases( + "openai.gpt-6.1-sol", + 48, + ExactExpected(spend=0.00079618, input_cost=0.00002178, output_cost=0.0007744, prompt_tokens=9, completion_tokens=64), + _mantle_gpt_usage(16, 5, 0), + ExactExpected(spend=0.00009922, input_cost=0.00003872, output_cost=0.0000605, prompt_tokens=16, completion_tokens=5), + 0, + ExactExpected(spend=0.000578622, input_cost=0.000385022, output_cost=0.0001936, prompt_tokens=3144, completion_tokens=16), + ExactExpected(spend=1.4594778, input_cost=1.452, output_cost=0.0074778, prompt_tokens=300000, completion_tokens=412), + ExactExpected(spend=0.5398778, input_cost=0.5324, output_cost=0.0074778, prompt_tokens=300000, completion_tokens=412), + ExactExpected(spend=0.6632252, input_cost=0.65824, output_cost=0.0049852, prompt_tokens=272000, completion_tokens=412), +) +_MANTLE_GPT_6_LUNA_CASES: Final = _mantle_gpt_cases( + "openai.gpt-6-luna", + 64, + ExactExpected(spend=0.000043428, input_cost=0.000001188, output_cost=0.00004224, prompt_tokens=9, completion_tokens=64), + _mantle_gpt_usage(16, 18, 11), + ExactExpected(spend=0.000013992, input_cost=0.000002112, output_cost=0.00001188, prompt_tokens=16, completion_tokens=18), + 16, + ExactExpected(spend=0.0000522984, input_cost=0.0000417384, output_cost=0.00001056, prompt_tokens=3144, completion_tokens=16), + ExactExpected(spend=0.07960788, input_cost=0.0792, output_cost=0.00040788, prompt_tokens=300000, completion_tokens=412), + ExactExpected(spend=0.03208788, input_cost=0.03168, output_cost=0.00040788, prompt_tokens=300000, completion_tokens=412), + ExactExpected(spend=0.03617592, input_cost=0.035904, output_cost=0.00027192, prompt_tokens=272000, completion_tokens=412), +) +_MANTLE_GPT_6_ASTRA_CASES: Final = _mantle_gpt_cases( + "openai.gpt-6-astra", + 29, + ExactExpected(spend=0.0047047, input_cost=0.0001287, output_cost=0.004576, prompt_tokens=9, completion_tokens=64), + _mantle_gpt_usage(16, 5, 0), + ExactExpected(spend=0.0005863, input_cost=0.0002288, output_cost=0.0003575, prompt_tokens=16, completion_tokens=5), + 0, + ExactExpected(spend=0.00566566, input_cost=0.00452166, output_cost=0.001144, prompt_tokens=3144, completion_tokens=16), + ExactExpected(spend=8.624187, input_cost=8.58, output_cost=0.044187, prompt_tokens=300000, completion_tokens=412), + ExactExpected(spend=3.476187, input_cost=3.432, output_cost=0.044187, prompt_tokens=300000, completion_tokens=412), + ExactExpected(spend=3.919058, input_cost=3.8896, output_cost=0.029458, prompt_tokens=272000, completion_tokens=412), +) +_MANTLE_GPT_56_TERRA_CASES: Final = _mantle_gpt_cases( + "openai.gpt-5.6-terra", + 33, + ExactExpected(spend=0.00121044, input_cost=0.00002772, output_cost=0.00118272, prompt_tokens=9, completion_tokens=64), + _mantle_gpt_usage(16, 17, 10), + ExactExpected(spend=0.00036344, input_cost=0.00004928, output_cost=0.00031416, prompt_tokens=16, completion_tokens=17), + 0, + ExactExpected(spend=0.001269576, input_cost=0.000973896, output_cost=0.00029568, prompt_tokens=3144, completion_tokens=16), + ExactExpected(spend=1.85942064, input_cost=1.848, output_cost=0.01142064, prompt_tokens=300000, completion_tokens=412), + ExactExpected(spend=0.75062064, input_cost=0.7392, output_cost=0.01142064, prompt_tokens=300000, completion_tokens=412), + ExactExpected(spend=0.84537376, input_cost=0.83776, output_cost=0.00761376, prompt_tokens=272000, completion_tokens=412), +) + +_MANTLE_CLAUDE_OPUS_MODEL: Final = "bedrock_mantle/anthropic.claude-opus-5-5" +_ADAPTIVE_USAGE: Final = { + "input_tokens": 24, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 18, + "output_tokens_details": {"thinking_tokens": 6}, + "service_tier": "standard", +} +_ADAPTIVE_EXPECTED: Final = ExactExpected( + spend=0.0007524, input_cost=0.0001584, output_cost=0.000594, prompt_tokens=24, completion_tokens=18 +) + +BEDROCK_MANTLE_CLAUDE_OPUS_5_5_INPUT_TEXT: Final = with_usage( + CLAUDE_OPUS_5_5_INPUT_TEXT, + "bedrock_mantle-anthropic.claude-opus-5-5-input_text", + _ADAPTIVE_USAGE, + _ADAPTIVE_EXPECTED, +).model_copy( + update={ + "model": _MANTLE_CLAUDE_OPUS_MODEL, + "request": { + **CLAUDE_OPUS_5_5_INPUT_TEXT.request, + "reasoning_effort": "medium", + "allowed_openai_params": ["reasoning_effort"], + }, + } +) + +_MANTLE_CLAUDE_STREAM_RESPONSE: Final = SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "type": "message_start", + "message": { + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-opus-5-5", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": { + "input_tokens": 24, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 5, + "service_tier": "standard", + }, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "thinking", "thinking": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "thinking_delta", "thinking": ""}}, + {"type": "content_block_stop", "index": 0}, + {"type": "content_block_start", "index": 1, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 1, "delta": {"type": "text_delta", "text": "4"}}, + {"type": "content_block_stop", "index": 1}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": { + "input_tokens": 24, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "output_tokens": 18, + "output_tokens_details": {"thinking_tokens": 6}, + }, + }, + {"type": "message_stop"}, + ), +) + +BEDROCK_MANTLE_CLAUDE_OPUS_5_5_STREAM: Final = CLAUDE_OPUS_5_5_STREAM.model_copy( + update={ + "name": "bedrock_mantle-anthropic.claude-opus-5-5-stream", + "model": _MANTLE_CLAUDE_OPUS_MODEL, + "request": { + **CLAUDE_OPUS_5_5_STREAM.request, + "reasoning_effort": "medium", + "allowed_openai_params": ["reasoning_effort"], + }, + "response": _MANTLE_CLAUDE_STREAM_RESPONSE, + "expected": _ADAPTIVE_EXPECTED, + } +) + +BEDROCK_MANTLE_CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT: Final = with_usage( + CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, + "bedrock_mantle-anthropic.claude-opus-5-5-messages_input_text", + _ADAPTIVE_USAGE, + _ADAPTIVE_EXPECTED, +).model_copy( + update={ + "model": _MANTLE_CLAUDE_OPUS_MODEL, + "request": { + **CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT.request, + "thinking": {"type": "adaptive"}, + "output_config": {"effort": "medium"}, + }, + } +) + +BEDROCK_MANTLE_CLAUDE_OPUS_5_5_MESSAGES_STREAM: Final = CLAUDE_OPUS_5_5_MESSAGES_STREAM.model_copy( + update={ + "name": "bedrock_mantle-anthropic.claude-opus-5-5-messages_stream", + "model": _MANTLE_CLAUDE_OPUS_MODEL, + "request": { + **CLAUDE_OPUS_5_5_MESSAGES_STREAM.request, + "thinking": {"type": "adaptive"}, + "output_config": {"effort": "medium"}, + }, + "response": _MANTLE_CLAUDE_STREAM_RESPONSE, + "expected": _ADAPTIVE_EXPECTED, + } +) + +BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CACHE_READ: Final = with_usage( + CLAUDE_OPUS_5_5_CACHE_READ, + "bedrock_mantle-anthropic.claude-opus-5-5-cache_read", + { + "input_tokens": 22, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 4203, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 47, + "output_tokens_details": {"thinking_tokens": 35}, + "service_tier": "standard", + }, + ExactExpected( + spend=0.00308319, input_cost=0.00153219, output_cost=0.001551, prompt_tokens=4225, completion_tokens=47 + ), +).model_copy(update={"model": _MANTLE_CLAUDE_OPUS_MODEL}) + +BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CACHE_WRITE_5M: Final = with_usage( + CLAUDE_OPUS_5_5_CACHE_WRITE_5M, + "bedrock_mantle-anthropic.claude-opus-5-5-cache_write_5m", + { + "input_tokens": 22, + "cache_creation_input_tokens": 4203, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 4203, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 44, + "output_tokens_details": {"thinking_tokens": 32}, + "service_tier": "standard", + }, + ExactExpected( + spend=0.03627195, input_cost=0.03481995, output_cost=0.001452, prompt_tokens=4225, completion_tokens=44 + ), +).model_copy(update={"model": _MANTLE_CLAUDE_OPUS_MODEL}) + +BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CACHE_WRITE_1H: Final = with_usage( + CLAUDE_OPUS_5_5_CACHE_WRITE_1H, + "bedrock_mantle-anthropic.claude-opus-5-5-cache_write_1h", + { + "input_tokens": 22, + "cache_creation_input_tokens": 9003, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 9003}, + "output_tokens": 12, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + }, + ExactExpected( + spend=0.1193808, input_cost=0.1189848, output_cost=0.000396, prompt_tokens=9025, completion_tokens=12 + ), +).model_copy(update={"model": _MANTLE_CLAUDE_OPUS_MODEL}) + +BEDROCK_MANTLE_CLAUDE_OPUS_5_5_THINKING: Final = with_usage( + CLAUDE_OPUS_5_5_THINKING, + "bedrock_mantle-anthropic.claude-opus-5-5-thinking", + _ADAPTIVE_USAGE, + _ADAPTIVE_EXPECTED, +).model_copy(update={"model": _MANTLE_CLAUDE_OPUS_MODEL}) + +BEDROCK_MANTLE_CLAUDE_OPUS_5_5_TOOL_CALL: Final = with_usage( + CLAUDE_OPUS_5_5_TOOL_CALL, + "bedrock_mantle-anthropic.claude-opus-5-5-tool_call", + { + "input_tokens": 376, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 12, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + }, + ExactExpected(spend=0.0028776, input_cost=0.0024816, output_cost=0.000396, prompt_tokens=376, completion_tokens=12), +).model_copy(update={"model": _MANTLE_CLAUDE_OPUS_MODEL}) + +BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CHAT_BRIDGED_JSON: Final = BEDROCK_MANTLE_CLAUDE_OPUS_5_5_INPUT_TEXT.model_copy( + update={ + "name": "bedrock_mantle-anthropic.claude-opus-5-5-chat_bridged_json", + "response": BEDROCK_MANTLE_CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT.response, + } +) + +BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( + _mantle_gpt_bridge(_MANTLE_GPT_61_SOL_CASES[0]), + _mantle_gpt_bridge(_MANTLE_GPT_6_LUNA_CASES[0]), + _mantle_gpt_bridge(_MANTLE_GPT_6_ASTRA_CASES[0]), + _mantle_gpt_bridge(_MANTLE_GPT_56_TERRA_CASES[0]), + BridgeCostParityTestCase( + native=BEDROCK_MANTLE_CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, + bridged=BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CHAT_BRIDGED_JSON, + ), +) + CASES: Final[tuple[CostTrackingTestCase, ...]] = ( BEDROCK_MANTLE_OPENAI_GPT_6_SOL_RESPONSES_JSON, BEDROCK_MANTLE_OPENAI_GPT_6_SOL_RESPONSES_STREAM, @@ -512,10 +941,35 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( BEDROCK_MANTLE_OPENAI_GPT_6_SOL_TIERED_INPUT_ABOVE_272K, BEDROCK_MANTLE_OPENAI_GPT_6_SOL_TIERED_CACHE_READ_ABOVE_272K, BEDROCK_MANTLE_OPENAI_GPT_6_SOL_TIERED_AT_272K_BOUNDARY, + *_MANTLE_GPT_61_SOL_CASES, + *_MANTLE_GPT_6_LUNA_CASES, + *_MANTLE_GPT_6_ASTRA_CASES, + *_MANTLE_GPT_56_TERRA_CASES, + BEDROCK_MANTLE_CLAUDE_OPUS_5_5_INPUT_TEXT, + BEDROCK_MANTLE_CLAUDE_OPUS_5_5_STREAM, + BEDROCK_MANTLE_CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, + BEDROCK_MANTLE_CLAUDE_OPUS_5_5_MESSAGES_STREAM, + BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CHAT_BRIDGED_JSON, + BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CACHE_READ, + BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CACHE_WRITE_5M, + BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CACHE_WRITE_1H, + BEDROCK_MANTLE_CLAUDE_OPUS_5_5_THINKING, + BEDROCK_MANTLE_CLAUDE_OPUS_5_5_TOOL_CALL, ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( StreamParityTestCase( plain=BEDROCK_MANTLE_OPENAI_GPT_6_SOL_RESPONSES_JSON, streamed=BEDROCK_MANTLE_OPENAI_GPT_6_SOL_RESPONSES_STREAM ), + StreamParityTestCase(plain=_MANTLE_GPT_61_SOL_CASES[0], streamed=_MANTLE_GPT_61_SOL_CASES[1]), + StreamParityTestCase(plain=_MANTLE_GPT_6_LUNA_CASES[0], streamed=_MANTLE_GPT_6_LUNA_CASES[1]), + StreamParityTestCase(plain=_MANTLE_GPT_6_ASTRA_CASES[0], streamed=_MANTLE_GPT_6_ASTRA_CASES[1]), + StreamParityTestCase(plain=_MANTLE_GPT_56_TERRA_CASES[0], streamed=_MANTLE_GPT_56_TERRA_CASES[1]), + StreamParityTestCase( + plain=BEDROCK_MANTLE_CLAUDE_OPUS_5_5_INPUT_TEXT, streamed=BEDROCK_MANTLE_CLAUDE_OPUS_5_5_STREAM + ), + StreamParityTestCase( + plain=BEDROCK_MANTLE_CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, + streamed=BEDROCK_MANTLE_CLAUDE_OPUS_5_5_MESSAGES_STREAM, + ), ) diff --git a/tests/integration/cost_calculation/cases/gemini.py b/tests/integration/cost_calculation/cases/gemini.py index 220ce2023be..bd64f4aeb11 100644 --- a/tests/integration/cost_calculation/cases/gemini.py +++ b/tests/integration/cost_calculation/cases/gemini.py @@ -87,6 +87,7 @@ GEMINI_GEMINI_3_1_PRO_INPUT_TEXT: Final = CostTrackingTestCase( "candidatesTokenCount": 412, "totalTokenCount": 2252, "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 1840}], + "serviceTier": "standard", }, "modelVersion": "gemini-3.1-pro", }, @@ -141,6 +142,7 @@ GEMINI_GEMINI_3_1_PRO_CACHE_READ: Final = CostTrackingTestCase( "totalTokenCount": 13308, "cachedContentTokenCount": 12288, "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 12928}], + "serviceTier": "standard", }, "modelVersion": "gemini-3.1-pro", }, @@ -386,6 +388,7 @@ GEMINI_GEMINI_3_1_PRO_REASONING: Final = CostTrackingTestCase( "thoughtsTokenCount": 3480, "totalTokenCount": 5280, "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 1240}], + "serviceTier": "standard", }, "modelVersion": "gemini-3.1-pro", }, @@ -1600,254 +1603,9 @@ GEMINI_GEMINI_3_8_FLASH_CACHE_READ: Final = CostTrackingTestCase( ), ) -GEMINI_GEMINI_3_8_FLASH_AUDIO_INPUT: Final = CostTrackingTestCase( - name="gemini-gemini-3.8-flash-audio_input", - covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="gemini/gemini-3.8-flash", - request={ - "model": "$MODEL", - "messages": [ - { - "role": "system", - "content": [ - { - "type": "text", - "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", - } - ], - }, - { - "role": "user", - "content": [ - { - "type": "text", - "text": "99e65f16c4b4 summarize the attached material in one line and name the city weather", - }, - { - "type": "input_audio", - "input_audio": { - "data": "UklGRmQGAABXQVZFZm10IBAAAAABAAEAQB8AAIA+AAACABAAZGF0YUAGAAAAAA8I4A85F+EdpCNYKNkrCy7eLkwuWSwTKZUk/x59GEERgQl4AWb5hPER6kbjVd1t2LLUQdIu0X/RM9M81oTa6t9E5mPtD/UP/SQFEw2fFI0bqSHEJrgqZi27LqouNi1pKlkmJSH0GvUTXgxpBFP8WPS27KflYt8U2ujV/dJp0TjRbNL71NPY1d3b47nqOPIh+jUCOQrwER0ZjB8JJW0plCxoLtou5i2VK/cnKSNOHZUWLw9VB0T/OPdv7yToj+Hi20nX5tPT0SDR09Hm00nX4tuP4STob+8490T/VQcvD5UWTh0pI/cnlSvmLdouaC6ULG0pCSWMHx0Z8BE5CjUCIfo48rnq2+PV3dPY+9Rs0jjRadH90ujVFNpi36fltuxY9FP8aQReDPUT9BolIVkmaSo2Laouuy5mLbgqxCapIY0bnxQTDSQFD/0P9WPtRObq34TaPNYz03/RLtFB0rLUbdhV3UbjEeqE8Wb5eAGBCUERfRj/HpUkEylZLEwu3i4LLtkrWCikI+EdORfgDw8IAADx9yDwx+gf4lzcqNcn1PXRItG00afT7dZr2wHhg+e/7n/2iP6aBnwO7xW6HKsikydOK78t0i6BLs0sxCl8JRYgvBmdEvEK8QLc+u3yYetz5FfePNlI1ZrSRdFW0crSl9Wn2dveDOUL7KLzl/utA6gLShNZGp4g7CUYKgMtly7ILpQtBSstJysiJRxHFcgN3wXL/cf1EO7j5nTg99qT1mzTmNEm0RrSa9QJ2NfcsuJr6dHwq/i8AMgIkRDcF3EeHiS3KBosLS7gLi0uGiy3KB4kcR7cF5EQyAi8AKv40fBr6bLi19wJ2GvUGtIm0ZjRbNOT1vfadODj5hDux/XL/d8FyA1HFSUcKyItJwUrlC3ILpcuAy0YKuwlniBZGkoTqAutA5f7ovML7Azl296n2ZfVytJW0UXRmtJI1TzZV95z5GHr7fLc+vEC8QqdErwZFiB8JcQpzSyBLtIuvy1OK5MnqyK6HO8VfA6aBoj+f/a/7oPnAeFr2+3Wp9O00SLR9dEn1KjXXNwf4sfoIPDx9wAADwjgDzkX4R2kI1go2SsLLt4uTC5ZLBMplST/Hn0YQRGBCXgBZvmE8RHqRuNV3W3YstRB0i7Rf9Ez0zzWhNrq30TmY+0P9Q/9JAUTDZ8UjRupIcQmuCpmLbsuqi42LWkqWSYlIfQa9RNeDGkEU/xY9Lbsp+Vi3xTa6NX90mnRONFs0vvU09jV3dvjueo48iH6NQI5CvARHRmMHwklbSmULGgu2i7mLZUr9ycpI04dlRYvD1UHRP8492/vJOiP4eLbSdfm09PRINHT0ebTSdfi24/hJOhv7zj3RP9VBy8PlRZOHSkj9yeVK+Yt2i5oLpQsbSkJJYwfHRnwETkKNQIh+jjyuerb49Xd09j71GzSONFp0f3S6NUU2mLfp+W27Fj0U/xpBF4M9RP0GiUhWSZpKjYtqi67LmYtuCrEJqkhjRufFBMNJAUP/Q/1Y+1E5urfhNo81jPTf9Eu0UHSstRt2FXdRuMR6oTxZvl4AYEJQRF9GP8elSQTKVksTC7eLgsu2StYKKQj4R05F+APDwgAAPH3IPDH6B/iXNyo1yfU9dEi0bTRp9Pt1mvbAeGD57/uf/aI/poGfA7vFbocqyKTJ04rvy3SLoEuzSzEKXwlFiC8GZ0S8QrxAtz67fJh63PkV9482UjVmtJF0VbRytKX1afZ294M5QvsovOX+60DqAtKE1kaniDsJRgqAy2XLsgulC0FKy0nKyIlHEcVyA3fBcv9x/UQ7uPmdOD32pPWbNOY0SbRGtJr1AnY19yy4mvp0fCr+LwAyAiRENwXcR4eJLcoGiwtLuAuLS4aLLcoHiRxHtwXkRDICLwAq/jR8GvpsuLX3AnYa9Qa0ibRmNFs05PW99p04OPmEO7H9cv93wXIDUcVJRwrIi0nBSuULcguly4DLRgq7CWeIFkaShOoC60Dl/ui8wvsDOXb3qfZl9XK0lbRRdGa0kjVPNlX3nPkYevt8tz68QLxCp0SvBkWIHwlxCnNLIEu0i6/LU4rkyerIroc7xV8DpoGiP5/9r/ug+cB4Wvb7dan07TRItH10SfUqNdc3B/ix+gg8PH3", - "format": "wav", - }, - }, - ], - }, - ], - "stream": False, - "modalities": ["text"], - "allowed_openai_params": ["modalities"], - }, - response=JsonResponse( - content_type="application/json", - body={ - "candidates": [ - { - "content": {"parts": [{"text": "scripted answer 99e65f16c4b4"}], "role": "model"}, - "finishReason": "STOP", - "index": 0, - } - ], - "usageMetadata": { - "promptTokenCount": 1546, - "candidatesTokenCount": 210, - "totalTokenCount": 1756, - "promptTokensDetails": [ - {"modality": "TEXT", "tokenCount": 96}, - {"modality": "AUDIO", "tokenCount": 1450}, - ], - }, - "modelVersion": "gemini-3.8-flash", - }, - ), - expected=ExactExpected( - spend=0.002128, input_cost=0.001498, output_cost=0.00063, prompt_tokens=1546, completion_tokens=210 - ), -) -GEMINI_GEMINI_3_8_FLASH_AUDIO_OUTPUT: Final = CostTrackingTestCase( - name="gemini-gemini-3.8-flash-audio_output", - covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="gemini/gemini-3.8-flash", - request={ - "model": "$MODEL", - "messages": [ - { - "role": "system", - "content": [ - { - "type": "text", - "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", - } - ], - }, - { - "role": "user", - "content": [ - { - "type": "text", - "text": "6fc6e4823e02 summarize the attached material in one line and name the city weather", - } - ], - }, - ], - "stream": False, - "modalities": ["text", "audio"], - "audio": {"voice": "alloy", "format": "pcm16"}, - "allowed_openai_params": ["modalities", "audio"], - }, - response=JsonResponse( - content_type="application/json", - body={ - "candidates": [ - { - "content": {"parts": [{"text": "scripted answer 6fc6e4823e02"}], "role": "model"}, - "finishReason": "STOP", - "index": 0, - } - ], - "usageMetadata": { - "promptTokenCount": 220, - "candidatesTokenCount": 1300, - "totalTokenCount": 1520, - "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 220}], - "candidatesTokensDetails": [ - {"modality": "TEXT", "tokenCount": 180}, - {"modality": "AUDIO", "tokenCount": 1120}, - ], - }, - "modelVersion": "gemini-3.8-flash", - }, - ), - expected=ExactExpected( - spend=0.00737, input_cost=0.00011, output_cost=0.00726, prompt_tokens=220, completion_tokens=1300 - ), -) -GEMINI_GEMINI_3_8_FLASH_IMAGE_INPUT: Final = CostTrackingTestCase( - name="gemini-gemini-3.8-flash-image_input", - covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="gemini/gemini-3.8-flash", - request={ - "model": "$MODEL", - "messages": [ - { - "role": "system", - "content": [ - { - "type": "text", - "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", - } - ], - }, - { - "role": "user", - "content": [ - { - "type": "text", - "text": "ca377dd90846 summarize the attached material in one line and name the city weather", - }, - { - "type": "image_url", - "image_url": { - "url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAACAAAAAgCAIAAAD8GO2jAAAMK0lEQVR4nAEgDN/zAM0HLNi+b59irEwJwoIG5+NVlKprNC9dCjpeSEL6tCj3YubiguXBZXx4w6lns2cR6zkGp8hgPXHUCeelTYe9wfcEQgJ6rx+pW3+GWJV430PkExZ66Nnc6zd2KDOBGnGnIwBzhiZIL2HGI3lifMEk1EYYPG5NnqGlpcz3LiFAwwS9/Goh5egSF1aINdSX+yErhrSeZWrPBkEWmgtZ9OYpQ58l2dRlT+yNSBn7QNa6ssjgEhxEGuYUp7jZKpIZrx7Oh1QAV1jeeB7zT4/0jccZhxGSWqLiJW9VRPJQuRZjnL/J8qO8F7vpSKdYNMGDc/cEMXKNBlAdehNaVHGe84Tdmm53hcOfr0ICjvEPuk0WzuaDIOumjneMSZ1+6qg8mAM8quAXAOyQPrjRNxDXsUwZZhYr07VMCinT0OP4yIEWDKvlaxGgReVKAJ1JpZx/Hlt+evT70yo3G94iWEhVavFwPnqJ87rKl0BT6+ohtOgz197MvB8QzMXpME/BwepPYkiRLpbBOAD37hU9bAWozdK3sPczg3okfSudzcFjAYtEIq5yw+1ZF9EYmBSexEP+IhnvUWC4BeDWZQiCjhF7/4syzu4C5kF9ETfrG+udK0232R+NA+uESn014bRJmPMfahYljDojL1UAbuaA0PuOGOzBBlCKXAkFNHIfvvZHQafMkl9qmqdFF4x3Em6WDuijSQXN6nFtM3UXbUGmmCF4Rcyl4YhifPspUXLdXZMIvPo9zAhTSlQFEi8m83sw5KxL0rWBzS9c4XAIAOaz3py4dTb7d9QaqDMLk0J87/15PZKvEaC6/hZ62sCtuFTywatjViHrBXTgOO5IJsiyYuzmaeQJKHmr11gnixR2rO7lqNEGs3shT+xK/lDUucFkigvA+a5C+itk/VV+1gDxc420UHpKhj31j0YnCZSFJebGz2/XSTyT6XfZhG4XNwZGIeUGCPKt0DX9lpB0RNN0yiPzQVJfa3LkZpRBN3RGgRpYc8WpHn5Q1wWpmnklpPHACv/N+kGz0Ki86g2GgvsAWloXy5pwfFtwZRYV9fMGU4Za35yfj4cddJuHfAyHSpYHVlGhNknUVQIBV9gOq7wwLZU3Pl5GJgSy4EK7f75iRVKD/B0ItpC0FBpwODhWP184ymnLgLGkK90WIVUY7hZtAEOu39BF2OsPDWrBGWN0esj6v3clX2z22gaLmrJHiwE4sXWUC8TLLtFp4uiS/VlbojLP9uiewb/vrTLBiFjYJ5rsFjuu/3XxEuiZ1QYyi9sfsFqPoiXjQjCA/jibX4aA1AAfp3GT1lykHtVMJmSxoW4Xvn3BXhXUdtWhIwP7NDO1HRP9UAlEUgKbd/iJBWS20DFBJQb2/UN/+CpSWi97Rta3+pe3H4kOr3qaV+g1Udgmupu6/cxmQqMP+DXd77Sy6a0AMxTVBR0DU4sVW/Vs2qPfnh3r+xlqt/3VIhyKQknN6xF5RIg4j7xsEpjsnKV/XhJNmN2sWekxom9lUCkuP3qgD25S7oCG6ZV3CrkUCm48s5hw+dUZANcGs4H6/PxIrSpkAC77CDPFF5hCvkfKWzerhufLBkq7AhZgeMKRnNboaP3m9qMh6+9Z3JEyaV8rfUacshIsMqwS/BI0uL9v9+fwcMQrbdwOstrkyWCPG61dXIAoEb9t2HnSdSlxy6FXa5+LhwCvCy1ANKkBHgVSx5hoE+LrBX47caskZapZ+MAsTFJhA3VxvHgKaWiuX4/vaG7TbOZdX7GRSzPz3yOeM4JeAuLqwOy6T0OBIKbnSm5b3A5+Y2j2cNYMiVmogx89QCIPRicA936Dj6rC2bDKBi8DmXw8dZvR170ELj4UjqD+VRopML3ww7ILyoVYi5P150ekt4QioS95PZdaHcPXSAD0GwVZe1t0K1qy2TGc7V2ySUwrZKwEnPRbLXAcl6praPJ8aVbkAJ1MPaI09ZDaZOT+npRQ4SFv1DK3halvT/cYVWNFxZy/HEwXatv4MtSF+5ymHEWqFC/kYwCJMTaYtzI7Ma9Q1rJ1WZtVBP36KFFdSj1G8Bw5b5ssozn/uHJxFO9g937ZtQBQvxvgA4h8rApfcpEHs+HfrYkWalhdEwiJ+fpmF/Um3ywbq7MF3kWROuUQa4H2rcVhq4WpdGuBte/A+QvHrGkq45kCcmvaWhAAslxCV5CWs7QlXig29URyTAgPh//Ji+IAJXC9fNNMdeirO7iPEEOemj9zZ8FMiAQAYKRF4o8E9glP+Jx+Vw9xU5oM40/363XW5T+Gd4xtwwwQoR3P38Wd0uURAeDvp2P53Wz6zxpyRmpc2iAwN5BujDYD2v+lpIn2AMUaEqKkrYP6sRheFZ0EY9Yt7r25gt9pITZVoPxxTMUDlHfdZgyMFfPLKLOtUk3gaiv88FD1Wd4AdI2pNnmYqAMKjqK35Ys3wVuBmgAbUPH6k4ae0ku7/qyRrkGHwRewnAAVZQgZavvBMFJ8cB4Ndb2bXEI2phv40TLtxqffTyNrTfbyrEdKKESwv/h/+pkLpi4bcaUZx0fBeRewm9oUozrsfuGNZcn0rKC8DdMUtqqXBWClU0bAHPTp2GD2c1FNwcwAFOXWyuK3ox0nFYLckD+pQW3ZJyKBKlCnanFXhVDRONHgupmm89QapXZ32FhAzncd/Xcy34ljiSee3UVHJtnq2YJJwQQVBNQKjuhoCv0YqzT/zVWuTbF5QtRm8I7cC5FQANj9TfJqwOxdaoY98K6RRZGzAezoRsF02Q3PwAowC6vYqUzHv85kEfb7BYqzoXGLm4fLxafznJZ+JhJdtsmV5qQiAxjnFQdrdTtL4KNZCXlvtdVYXxLGFN9os7WJoU2kLQB5RTU1rvAlYHfbTZZTTYFKXxRHAizqcSNldKkms7p4RYVAWlUXIGdWTb4kym3qAF4clO+3Njf9F3zxl1b73YfH8pXbnk9fIQnHRowK9vQFRcfD8ikeIUAmZrhe+3cOXZUAt7EeSjZgZFxWFhFll/ZC/Y93aYwMJjAhvbgcS/zZaR0ucmJ663gDu7xgXT3tgQrvb4daDFM10U5HZ/kt2r9m26me5oXdgom65hEkEgfTe03ZFzZnPz7n1fzuGVThCppMAAAyK2pIlNAb40xyGcWNko8VyJ3oghsue2leWHkJ6UpWCdZB13x+QcxkLOr5a8X0zx6N+VcXwNMfGGqlelK0sh1MrRizk4drf0prMW3Q264J3IVJTZ3tx1ULpBi8CnnmrwBeYvfkI5rZJ2W6cOv3bKMqoCpyoNo+gpCZcSVgB8hdzlfMfPl67dL9/IqNo2mBQBaVitoQ4y7KzqHnrBWSxLIyznP3v/VNlzgF7qcPaQiDYhyGAgLfihsZE5nQZLcedY4Apm3dhEOyMUpdzNNyViCP6ykfFn35Ss3JDUS7lfZV47bpxeatoe2dNzU5TGxY9GoVYvhKXmP0WjgOpwv91XPF/x1ty0JRgNGVMwnwxk4PpSpDgGRu/k4jMev9x1vlc3EdAIz+WBCWU23SZoIT3WfyVNpXi7azEQkSB7e1K+fS25uGiNScc4/Nn8BVsUqtwtGhM5nncB7tH0F/k0ngZeGmqEGNJw01v2UYpLQoxwqhCz7/X5it8liVJMxk2kQ3iCGYnwDwfHcl7gLJGiAqrjLBOlp17bPGYIGTlFw+TOCLsGS103TZvSbyptAazsca/TY3o5hZLqX7SBvnfL6s/SU7BVyQY4PpZwHLa6PI2w+uFmJZbBsaruCVPW+FL+zSfvPO560ABJLD6nKsReNBwoxShnNi1XPSmX+C7Ij/ik2nfzubk6IFlomrI0bs386ztTMDifWKNO0ZnC+Oau6ZgwE4PwF8AqKtGFd5yqG0aMFq9rvwPb8DZCXmUHN5Uc40cNG5HWNSAK1yFAkCcG2GPV/7JYEi3t9ZA8NBeh9Dmvsjv8jiaeknPCij14x6BmGOFzyVNqRcS8x517fM39S0cC6bzu8wbnh96fwQ9z/JzCEsqxV9KjuES+xv3rYkX+r3JxcK7p78EACPip85gRBoIU7cZgrwLepMArrpA9YvR8Icag3YILzWcgXVoVRTs6Xc2PsiKasZ98y5CBkmP9Ri3lwZ9acgvPJ1fw729yDzDl/0qEeBOf7ZYZ7dreaPeaGHIymfgKAwmwhnjvAGaqLWhgAAAABJRU5ErkJggg==", - "detail": "high", - }, - }, - ], - }, - ], - "stream": False, - "allowed_openai_params": [], - }, - response=JsonResponse( - content_type="application/json", - body={ - "candidates": [ - { - "content": {"parts": [{"text": "scripted answer ca377dd90846"}], "role": "model"}, - "finishReason": "STOP", - "index": 0, - } - ], - "usageMetadata": { - "promptTokenCount": 2116, - "candidatesTokenCount": 240, - "totalTokenCount": 2356, - "promptTokensDetails": [ - {"modality": "TEXT", "tokenCount": 310}, - {"modality": "IMAGE", "tokenCount": 1806}, - ], - }, - "modelVersion": "gemini-3.8-flash", - }, - ), - expected=ExactExpected( - spend=0.0018683, input_cost=0.0011483, output_cost=0.00072, prompt_tokens=2116, completion_tokens=240 - ), -) -GEMINI_GEMINI_3_8_FLASH_VIDEO_INPUT: Final = CostTrackingTestCase( - name="gemini-gemini-3.8-flash-video_input", - covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="gemini/gemini-3.8-flash", - request={ - "model": "$MODEL", - "messages": [ - { - "role": "system", - "content": [ - { - "type": "text", - "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", - } - ], - }, - { - "role": "user", - "content": [ - { - "type": "text", - "text": "030071a5c80f summarize the attached material in one line and name the city weather", - }, - { - "type": "file", - "file": { - "file_data": "data:video/mp4;base64,AAAAGGZ0eXBpc29tAAACAGlzb21pc282AAACCG1kYXQNFBsiKTA3PkVMU1phaG92fYSLkpmgp661vMPK0djf5u30+wIJEBceJSwzOkFIT1ZdZGtyeYCHjpWco6qxuL/GzdTb4unw9/4FDBMaISgvNj1ES1JZYGdudXyDipGYn6attLvCydDX3uXs8/oBCA8WHSQrMjlAR05VXGNqcXh/ho2Um6KpsLe+xczT2uHo7/b9BAsSGSAnLjU8Q0pRWF9mbXR7gomQl56lrLO6wcjP1t3k6/L5AAcOFRwjKjE4P0ZNVFtiaXB3foWMk5qhqK+2vcTL0tng5+71/AMKERgfJi00O0JJUFdeZWxzeoGIj5adpKuyucDHztXc4+rx+P8GDRQbIikwNz5FTFNaYWhvdn2Ei5KZoKeutbzDytHY3+bt9PsCCRAXHiUsMzpBSE9WXWRrcnmAh46VnKOqsbi/xs3U2+Lp8Pf+BQwTGiEoLzY9REtSWWBnbnV8g4qRmJ+mrbS7wsnQ197l7PP6AQgPFh0kKzI5QEdOVVxjanF4f4aNlJuiqbC3vsXM09rh6O/2/QQLEhkgJy41PENKUVhfZm10e4KJkJeepayzusHIz9bd5Ovy+QAHDhUcIyoxOD9GTVRbYmlwd36FjJOaoaivtr3Ey9LZ4Ofu9fwDChEYHyYtNDtCSVBXXmVsc3qBiI+WnaSrsrnAx87V3OPq8fj/Bg==", - "format": "mp4", - }, - }, - ], - }, - ], - "stream": False, - "allowed_openai_params": [], - }, - response=JsonResponse( - content_type="application/json", - body={ - "candidates": [ - { - "content": {"parts": [{"text": "scripted answer 030071a5c80f"}], "role": "model"}, - "finishReason": "STOP", - "index": 0, - } - ], - "usageMetadata": { - "promptTokenCount": 8060, - "candidatesTokenCount": 300, - "totalTokenCount": 8360, - "promptTokensDetails": [ - {"modality": "TEXT", "tokenCount": 140}, - {"modality": "VIDEO", "tokenCount": 7920}, - ], - }, - "modelVersion": "gemini-3.8-flash", - }, - ), - expected=ExactExpected( - spend=0.005722, input_cost=0.004822, output_cost=0.0009, prompt_tokens=8060, completion_tokens=300 - ), -) GEMINI_GEMINI_3_8_FLASH_REASONING: Final = CostTrackingTestCase( name="gemini-gemini-3.8-flash-reasoning", @@ -2316,60 +2074,6 @@ GEMINI_GEMINI_3_8_FLASH_STREAM_NO_USAGE_TOOL_CALL: Final = CostTrackingTestCase( expected=RecountExpected(recount=RecountRates(input_cost_per_token=5e-07, output_cost_per_token=3e-06)), ) -GEMINI_GEMINI_3_8_FLASH_STREAM_NO_USAGE_IMAGE_INPUT: Final = CostTrackingTestCase( - name="gemini-gemini-3.8-flash-stream_no_usage_image_input", - covers="quota_management.spend_tracking.scripted_wire.logs_cost", - model="gemini/gemini-3.8-flash", - request={ - "model": "$MODEL", - "messages": [ - { - "role": "system", - "content": [ - { - "type": "text", - "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", - } - ], - }, - { - "role": "user", - "content": [ - { - "type": "text", - "text": "5fca57bc9ae9 summarize the attached material in one line and name the city weather", - }, - { - "type": "image_url", - "image_url": { - "url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAACAAAAAgCAIAAAD8GO2jAAAMK0lEQVR4nAEgDN/zAM0HLNi+b59irEwJwoIG5+NVlKprNC9dCjpeSEL6tCj3YubiguXBZXx4w6lns2cR6zkGp8hgPXHUCeelTYe9wfcEQgJ6rx+pW3+GWJV430PkExZ66Nnc6zd2KDOBGnGnIwBzhiZIL2HGI3lifMEk1EYYPG5NnqGlpcz3LiFAwwS9/Goh5egSF1aINdSX+yErhrSeZWrPBkEWmgtZ9OYpQ58l2dRlT+yNSBn7QNa6ssjgEhxEGuYUp7jZKpIZrx7Oh1QAV1jeeB7zT4/0jccZhxGSWqLiJW9VRPJQuRZjnL/J8qO8F7vpSKdYNMGDc/cEMXKNBlAdehNaVHGe84Tdmm53hcOfr0ICjvEPuk0WzuaDIOumjneMSZ1+6qg8mAM8quAXAOyQPrjRNxDXsUwZZhYr07VMCinT0OP4yIEWDKvlaxGgReVKAJ1JpZx/Hlt+evT70yo3G94iWEhVavFwPnqJ87rKl0BT6+ohtOgz197MvB8QzMXpME/BwepPYkiRLpbBOAD37hU9bAWozdK3sPczg3okfSudzcFjAYtEIq5yw+1ZF9EYmBSexEP+IhnvUWC4BeDWZQiCjhF7/4syzu4C5kF9ETfrG+udK0232R+NA+uESn014bRJmPMfahYljDojL1UAbuaA0PuOGOzBBlCKXAkFNHIfvvZHQafMkl9qmqdFF4x3Em6WDuijSQXN6nFtM3UXbUGmmCF4Rcyl4YhifPspUXLdXZMIvPo9zAhTSlQFEi8m83sw5KxL0rWBzS9c4XAIAOaz3py4dTb7d9QaqDMLk0J87/15PZKvEaC6/hZ62sCtuFTywatjViHrBXTgOO5IJsiyYuzmaeQJKHmr11gnixR2rO7lqNEGs3shT+xK/lDUucFkigvA+a5C+itk/VV+1gDxc420UHpKhj31j0YnCZSFJebGz2/XSTyT6XfZhG4XNwZGIeUGCPKt0DX9lpB0RNN0yiPzQVJfa3LkZpRBN3RGgRpYc8WpHn5Q1wWpmnklpPHACv/N+kGz0Ki86g2GgvsAWloXy5pwfFtwZRYV9fMGU4Za35yfj4cddJuHfAyHSpYHVlGhNknUVQIBV9gOq7wwLZU3Pl5GJgSy4EK7f75iRVKD/B0ItpC0FBpwODhWP184ymnLgLGkK90WIVUY7hZtAEOu39BF2OsPDWrBGWN0esj6v3clX2z22gaLmrJHiwE4sXWUC8TLLtFp4uiS/VlbojLP9uiewb/vrTLBiFjYJ5rsFjuu/3XxEuiZ1QYyi9sfsFqPoiXjQjCA/jibX4aA1AAfp3GT1lykHtVMJmSxoW4Xvn3BXhXUdtWhIwP7NDO1HRP9UAlEUgKbd/iJBWS20DFBJQb2/UN/+CpSWi97Rta3+pe3H4kOr3qaV+g1Udgmupu6/cxmQqMP+DXd77Sy6a0AMxTVBR0DU4sVW/Vs2qPfnh3r+xlqt/3VIhyKQknN6xF5RIg4j7xsEpjsnKV/XhJNmN2sWekxom9lUCkuP3qgD25S7oCG6ZV3CrkUCm48s5hw+dUZANcGs4H6/PxIrSpkAC77CDPFF5hCvkfKWzerhufLBkq7AhZgeMKRnNboaP3m9qMh6+9Z3JEyaV8rfUacshIsMqwS/BI0uL9v9+fwcMQrbdwOstrkyWCPG61dXIAoEb9t2HnSdSlxy6FXa5+LhwCvCy1ANKkBHgVSx5hoE+LrBX47caskZapZ+MAsTFJhA3VxvHgKaWiuX4/vaG7TbOZdX7GRSzPz3yOeM4JeAuLqwOy6T0OBIKbnSm5b3A5+Y2j2cNYMiVmogx89QCIPRicA936Dj6rC2bDKBi8DmXw8dZvR170ELj4UjqD+VRopML3ww7ILyoVYi5P150ekt4QioS95PZdaHcPXSAD0GwVZe1t0K1qy2TGc7V2ySUwrZKwEnPRbLXAcl6praPJ8aVbkAJ1MPaI09ZDaZOT+npRQ4SFv1DK3halvT/cYVWNFxZy/HEwXatv4MtSF+5ymHEWqFC/kYwCJMTaYtzI7Ma9Q1rJ1WZtVBP36KFFdSj1G8Bw5b5ssozn/uHJxFO9g937ZtQBQvxvgA4h8rApfcpEHs+HfrYkWalhdEwiJ+fpmF/Um3ywbq7MF3kWROuUQa4H2rcVhq4WpdGuBte/A+QvHrGkq45kCcmvaWhAAslxCV5CWs7QlXig29URyTAgPh//Ji+IAJXC9fNNMdeirO7iPEEOemj9zZ8FMiAQAYKRF4o8E9glP+Jx+Vw9xU5oM40/363XW5T+Gd4xtwwwQoR3P38Wd0uURAeDvp2P53Wz6zxpyRmpc2iAwN5BujDYD2v+lpIn2AMUaEqKkrYP6sRheFZ0EY9Yt7r25gt9pITZVoPxxTMUDlHfdZgyMFfPLKLOtUk3gaiv88FD1Wd4AdI2pNnmYqAMKjqK35Ys3wVuBmgAbUPH6k4ae0ku7/qyRrkGHwRewnAAVZQgZavvBMFJ8cB4Ndb2bXEI2phv40TLtxqffTyNrTfbyrEdKKESwv/h/+pkLpi4bcaUZx0fBeRewm9oUozrsfuGNZcn0rKC8DdMUtqqXBWClU0bAHPTp2GD2c1FNwcwAFOXWyuK3ox0nFYLckD+pQW3ZJyKBKlCnanFXhVDRONHgupmm89QapXZ32FhAzncd/Xcy34ljiSee3UVHJtnq2YJJwQQVBNQKjuhoCv0YqzT/zVWuTbF5QtRm8I7cC5FQANj9TfJqwOxdaoY98K6RRZGzAezoRsF02Q3PwAowC6vYqUzHv85kEfb7BYqzoXGLm4fLxafznJZ+JhJdtsmV5qQiAxjnFQdrdTtL4KNZCXlvtdVYXxLGFN9os7WJoU2kLQB5RTU1rvAlYHfbTZZTTYFKXxRHAizqcSNldKkms7p4RYVAWlUXIGdWTb4kym3qAF4clO+3Njf9F3zxl1b73YfH8pXbnk9fIQnHRowK9vQFRcfD8ikeIUAmZrhe+3cOXZUAt7EeSjZgZFxWFhFll/ZC/Y93aYwMJjAhvbgcS/zZaR0ucmJ663gDu7xgXT3tgQrvb4daDFM10U5HZ/kt2r9m26me5oXdgom65hEkEgfTe03ZFzZnPz7n1fzuGVThCppMAAAyK2pIlNAb40xyGcWNko8VyJ3oghsue2leWHkJ6UpWCdZB13x+QcxkLOr5a8X0zx6N+VcXwNMfGGqlelK0sh1MrRizk4drf0prMW3Q264J3IVJTZ3tx1ULpBi8CnnmrwBeYvfkI5rZJ2W6cOv3bKMqoCpyoNo+gpCZcSVgB8hdzlfMfPl67dL9/IqNo2mBQBaVitoQ4y7KzqHnrBWSxLIyznP3v/VNlzgF7qcPaQiDYhyGAgLfihsZE5nQZLcedY4Apm3dhEOyMUpdzNNyViCP6ykfFn35Ss3JDUS7lfZV47bpxeatoe2dNzU5TGxY9GoVYvhKXmP0WjgOpwv91XPF/x1ty0JRgNGVMwnwxk4PpSpDgGRu/k4jMev9x1vlc3EdAIz+WBCWU23SZoIT3WfyVNpXi7azEQkSB7e1K+fS25uGiNScc4/Nn8BVsUqtwtGhM5nncB7tH0F/k0ngZeGmqEGNJw01v2UYpLQoxwqhCz7/X5it8liVJMxk2kQ3iCGYnwDwfHcl7gLJGiAqrjLBOlp17bPGYIGTlFw+TOCLsGS103TZvSbyptAazsca/TY3o5hZLqX7SBvnfL6s/SU7BVyQY4PpZwHLa6PI2w+uFmJZbBsaruCVPW+FL+zSfvPO560ABJLD6nKsReNBwoxShnNi1XPSmX+C7Ij/ik2nfzubk6IFlomrI0bs386ztTMDifWKNO0ZnC+Oau6ZgwE4PwF8AqKtGFd5yqG0aMFq9rvwPb8DZCXmUHN5Uc40cNG5HWNSAK1yFAkCcG2GPV/7JYEi3t9ZA8NBeh9Dmvsjv8jiaeknPCij14x6BmGOFzyVNqRcS8x517fM39S0cC6bzu8wbnh96fwQ9z/JzCEsqxV9KjuES+xv3rYkX+r3JxcK7p78EACPip85gRBoIU7cZgrwLepMArrpA9YvR8Icag3YILzWcgXVoVRTs6Xc2PsiKasZ98y5CBkmP9Ri3lwZ9acgvPJ1fw729yDzDl/0qEeBOf7ZYZ7dreaPeaGHIymfgKAwmwhnjvAGaqLWhgAAAABJRU5ErkJggg==", - "detail": "high", - }, - }, - ], - }, - ], - "stream": True, - "stream_options": {"include_usage": True}, - "allowed_openai_params": [], - }, - response=SseResponse( - content_type="text/event-stream", - frames=sse_frames( - { - "candidates": [ - { - "content": {"parts": [{"text": "scripted answer 5fca57bc9ae9"}], "role": "model"}, - "finishReason": "STOP", - "index": 0, - } - ], - "modelVersion": "gemini-3.8-flash", - } - ), - ), - expected=RecountExpected(recount=RecountRates(input_cost_per_token=5e-07, output_cost_per_token=3e-06)), -) GEMINI_GEMINI_3_8_FLASH_PROMPT_BLOCKED: Final = CostTrackingTestCase( name="gemini-gemini-3.8-flash-prompt_blocked", @@ -3641,10 +3345,6 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GEMINI_GEMINI_3_1_PRO_STREAM_FULL_USAGE, GEMINI_GEMINI_3_8_FLASH_INPUT_TEXT, GEMINI_GEMINI_3_8_FLASH_CACHE_READ, - GEMINI_GEMINI_3_8_FLASH_AUDIO_INPUT, - GEMINI_GEMINI_3_8_FLASH_AUDIO_OUTPUT, - GEMINI_GEMINI_3_8_FLASH_IMAGE_INPUT, - GEMINI_GEMINI_3_8_FLASH_VIDEO_INPUT, GEMINI_GEMINI_3_8_FLASH_REASONING, GEMINI_GEMINI_3_8_FLASH_SERVICE_TIER_FLEX, GEMINI_GEMINI_3_8_FLASH_SERVICE_TIER_PRIORITY, @@ -3653,7 +3353,6 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GEMINI_GEMINI_3_8_FLASH_STREAM, GEMINI_GEMINI_3_8_FLASH_STREAM_NO_USAGE, GEMINI_GEMINI_3_8_FLASH_STREAM_NO_USAGE_TOOL_CALL, - GEMINI_GEMINI_3_8_FLASH_STREAM_NO_USAGE_IMAGE_INPUT, GEMINI_GEMINI_3_8_FLASH_PROMPT_BLOCKED, GEMINI_GEMINI_3_8_FLASH_STREAM_PROMPT_BLOCKED, GEMINI_GEMINI_3_8_FLASH_RESPONSE_MODEL_OVERRIDE, diff --git a/tests/integration/cost_calculation/cases/vertex_ai.py b/tests/integration/cost_calculation/cases/vertex_ai.py index e19e377ae93..761078ef593 100644 --- a/tests/integration/cost_calculation/cases/vertex_ai.py +++ b/tests/integration/cost_calculation/cases/vertex_ai.py @@ -248,7 +248,9 @@ GEMINI_3_1_PRO_INPUT_TEXT: Final = CostTrackingTestCase( "promptTokenCount": 1840, "candidatesTokenCount": 412, "totalTokenCount": 2252, + "trafficType": "ON_DEMAND", "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 1840}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 412}], }, "modelVersion": "gemini-3.1-pro", }, @@ -302,7 +304,9 @@ GEMINI_3_1_PRO_CACHE_READ: Final = CostTrackingTestCase( "candidatesTokenCount": 380, "totalTokenCount": 13308, "cachedContentTokenCount": 12288, + "trafficType": "ON_DEMAND", "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 12928}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 380}], }, "modelVersion": "gemini-3.1-pro", }, @@ -935,7 +939,9 @@ GEMINI_3_1_PRO_STREAM: Final = CostTrackingTestCase( "promptTokenCount": 1840, "candidatesTokenCount": 412, "totalTokenCount": 2252, + "trafficType": "ON_DEMAND", "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 1840}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 412}], }, "modelVersion": "gemini-3.1-pro", }, @@ -1667,9 +1673,9 @@ GEMINI_3_1_PRO_PASSTHROUGH_STREAM_GENERATE_CONTENT_PRICED_VIA_VERTEX_KEY: Final # gemini-3.8-flash GEMINI_3_8_FLASH_INPUT_TEXT: Final = CostTrackingTestCase( - name="gemini-3.8-flash-input_text", + name="vertex_ai/gemini-3.8-flash-input_text", covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="gemini-3.8-flash", + model="vertex_ai/gemini-3.8-flash", request={ "model": "$MODEL", "messages": [ @@ -1720,9 +1726,9 @@ GEMINI_3_8_FLASH_INPUT_TEXT: Final = CostTrackingTestCase( ) GEMINI_3_8_FLASH_CACHE_READ: Final = CostTrackingTestCase( - name="gemini-3.8-flash-cache_read", + name="vertex_ai/gemini-3.8-flash-cache_read", covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="gemini-3.8-flash", + model="vertex_ai/gemini-3.8-flash", request={ "model": "$MODEL", "messages": [ @@ -1773,196 +1779,13 @@ GEMINI_3_8_FLASH_CACHE_READ: Final = CostTrackingTestCase( ), ) -GEMINI_3_8_FLASH_AUDIO_INPUT: Final = CostTrackingTestCase( - name="gemini-3.8-flash-audio_input", - covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="gemini-3.8-flash", - request={ - "model": "$MODEL", - "messages": [ - { - "role": "system", - "content": [ - { - "type": "text", - "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", - } - ], - }, - { - "role": "user", - "content": [ - { - "type": "text", - "text": "3e33892c4f9f summarize the attached material in one line and name the city weather", - }, - { - "type": "input_audio", - "input_audio": { - "data": "UklGRmQGAABXQVZFZm10IBAAAAABAAEAQB8AAIA+AAACABAAZGF0YUAGAAAAAA8I4A85F+EdpCNYKNkrCy7eLkwuWSwTKZUk/x59GEERgQl4AWb5hPER6kbjVd1t2LLUQdIu0X/RM9M81oTa6t9E5mPtD/UP/SQFEw2fFI0bqSHEJrgqZi27LqouNi1pKlkmJSH0GvUTXgxpBFP8WPS27KflYt8U2ujV/dJp0TjRbNL71NPY1d3b47nqOPIh+jUCOQrwER0ZjB8JJW0plCxoLtou5i2VK/cnKSNOHZUWLw9VB0T/OPdv7yToj+Hi20nX5tPT0SDR09Hm00nX4tuP4STob+8490T/VQcvD5UWTh0pI/cnlSvmLdouaC6ULG0pCSWMHx0Z8BE5CjUCIfo48rnq2+PV3dPY+9Rs0jjRadH90ujVFNpi36fltuxY9FP8aQReDPUT9BolIVkmaSo2Laouuy5mLbgqxCapIY0bnxQTDSQFD/0P9WPtRObq34TaPNYz03/RLtFB0rLUbdhV3UbjEeqE8Wb5eAGBCUERfRj/HpUkEylZLEwu3i4LLtkrWCikI+EdORfgDw8IAADx9yDwx+gf4lzcqNcn1PXRItG00afT7dZr2wHhg+e/7n/2iP6aBnwO7xW6HKsikydOK78t0i6BLs0sxCl8JRYgvBmdEvEK8QLc+u3yYetz5FfePNlI1ZrSRdFW0crSl9Wn2dveDOUL7KLzl/utA6gLShNZGp4g7CUYKgMtly7ILpQtBSstJysiJRxHFcgN3wXL/cf1EO7j5nTg99qT1mzTmNEm0RrSa9QJ2NfcsuJr6dHwq/i8AMgIkRDcF3EeHiS3KBosLS7gLi0uGiy3KB4kcR7cF5EQyAi8AKv40fBr6bLi19wJ2GvUGtIm0ZjRbNOT1vfadODj5hDux/XL/d8FyA1HFSUcKyItJwUrlC3ILpcuAy0YKuwlniBZGkoTqAutA5f7ovML7Azl296n2ZfVytJW0UXRmtJI1TzZV95z5GHr7fLc+vEC8QqdErwZFiB8JcQpzSyBLtIuvy1OK5MnqyK6HO8VfA6aBoj+f/a/7oPnAeFr2+3Wp9O00SLR9dEn1KjXXNwf4sfoIPDx9wAADwjgDzkX4R2kI1go2SsLLt4uTC5ZLBMplST/Hn0YQRGBCXgBZvmE8RHqRuNV3W3YstRB0i7Rf9Ez0zzWhNrq30TmY+0P9Q/9JAUTDZ8UjRupIcQmuCpmLbsuqi42LWkqWSYlIfQa9RNeDGkEU/xY9Lbsp+Vi3xTa6NX90mnRONFs0vvU09jV3dvjueo48iH6NQI5CvARHRmMHwklbSmULGgu2i7mLZUr9ycpI04dlRYvD1UHRP8492/vJOiP4eLbSdfm09PRINHT0ebTSdfi24/hJOhv7zj3RP9VBy8PlRZOHSkj9yeVK+Yt2i5oLpQsbSkJJYwfHRnwETkKNQIh+jjyuerb49Xd09j71GzSONFp0f3S6NUU2mLfp+W27Fj0U/xpBF4M9RP0GiUhWSZpKjYtqi67LmYtuCrEJqkhjRufFBMNJAUP/Q/1Y+1E5urfhNo81jPTf9Eu0UHSstRt2FXdRuMR6oTxZvl4AYEJQRF9GP8elSQTKVksTC7eLgsu2StYKKQj4R05F+APDwgAAPH3IPDH6B/iXNyo1yfU9dEi0bTRp9Pt1mvbAeGD57/uf/aI/poGfA7vFbocqyKTJ04rvy3SLoEuzSzEKXwlFiC8GZ0S8QrxAtz67fJh63PkV9482UjVmtJF0VbRytKX1afZ294M5QvsovOX+60DqAtKE1kaniDsJRgqAy2XLsgulC0FKy0nKyIlHEcVyA3fBcv9x/UQ7uPmdOD32pPWbNOY0SbRGtJr1AnY19yy4mvp0fCr+LwAyAiRENwXcR4eJLcoGiwtLuAuLS4aLLcoHiRxHtwXkRDICLwAq/jR8GvpsuLX3AnYa9Qa0ibRmNFs05PW99p04OPmEO7H9cv93wXIDUcVJRwrIi0nBSuULcguly4DLRgq7CWeIFkaShOoC60Dl/ui8wvsDOXb3qfZl9XK0lbRRdGa0kjVPNlX3nPkYevt8tz68QLxCp0SvBkWIHwlxCnNLIEu0i6/LU4rkyerIroc7xV8DpoGiP5/9r/ug+cB4Wvb7dan07TRItH10SfUqNdc3B/ix+gg8PH3", - "format": "wav", - }, - }, - ], - }, - ], - "stream": False, - "modalities": ["text"], - "allowed_openai_params": ["modalities"], - }, - response=JsonResponse( - content_type="application/json", - body={ - "candidates": [ - { - "content": {"parts": [{"text": "scripted answer 3e33892c4f9f"}], "role": "model"}, - "finishReason": "STOP", - "index": 0, - } - ], - "usageMetadata": { - "promptTokenCount": 1546, - "candidatesTokenCount": 210, - "totalTokenCount": 1756, - "promptTokensDetails": [ - {"modality": "TEXT", "tokenCount": 96}, - {"modality": "AUDIO", "tokenCount": 1450}, - ], - }, - "modelVersion": "gemini-3.8-flash", - }, - ), - expected=ExactExpected( - spend=0.00221312, input_cost=0.00155792, output_cost=0.0006552, prompt_tokens=1546, completion_tokens=210 - ), -) -GEMINI_3_8_FLASH_AUDIO_OUTPUT: Final = CostTrackingTestCase( - name="gemini-3.8-flash-audio_output", - covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="gemini-3.8-flash", - request={ - "model": "$MODEL", - "messages": [ - { - "role": "system", - "content": [ - { - "type": "text", - "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", - } - ], - }, - { - "role": "user", - "content": [ - { - "type": "text", - "text": "30bf1c0de6fe summarize the attached material in one line and name the city weather", - } - ], - }, - ], - "stream": False, - "modalities": ["text", "audio"], - "audio": {"voice": "alloy", "format": "pcm16"}, - "allowed_openai_params": ["modalities", "audio"], - }, - response=JsonResponse( - content_type="application/json", - body={ - "candidates": [ - { - "content": {"parts": [{"text": "scripted answer 30bf1c0de6fe"}], "role": "model"}, - "finishReason": "STOP", - "index": 0, - } - ], - "usageMetadata": { - "promptTokenCount": 220, - "candidatesTokenCount": 1300, - "totalTokenCount": 1520, - "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 220}], - "candidatesTokensDetails": [ - {"modality": "TEXT", "tokenCount": 180}, - {"modality": "AUDIO", "tokenCount": 1120}, - ], - }, - "modelVersion": "gemini-3.8-flash", - }, - ), - expected=ExactExpected( - spend=0.0076648, input_cost=0.0001144, output_cost=0.0075504, prompt_tokens=220, completion_tokens=1300 - ), -) -GEMINI_3_8_FLASH_VIDEO_INPUT: Final = CostTrackingTestCase( - name="gemini-3.8-flash-video_input", - covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="gemini-3.8-flash", - request={ - "model": "$MODEL", - "messages": [ - { - "role": "system", - "content": [ - { - "type": "text", - "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", - } - ], - }, - { - "role": "user", - "content": [ - { - "type": "text", - "text": "5f5da5957185 summarize the attached material in one line and name the city weather", - }, - { - "type": "file", - "file": { - "file_data": "data:video/mp4;base64,AAAAGGZ0eXBpc29tAAACAGlzb21pc282AAACCG1kYXQNFBsiKTA3PkVMU1phaG92fYSLkpmgp661vMPK0djf5u30+wIJEBceJSwzOkFIT1ZdZGtyeYCHjpWco6qxuL/GzdTb4unw9/4FDBMaISgvNj1ES1JZYGdudXyDipGYn6attLvCydDX3uXs8/oBCA8WHSQrMjlAR05VXGNqcXh/ho2Um6KpsLe+xczT2uHo7/b9BAsSGSAnLjU8Q0pRWF9mbXR7gomQl56lrLO6wcjP1t3k6/L5AAcOFRwjKjE4P0ZNVFtiaXB3foWMk5qhqK+2vcTL0tng5+71/AMKERgfJi00O0JJUFdeZWxzeoGIj5adpKuyucDHztXc4+rx+P8GDRQbIikwNz5FTFNaYWhvdn2Ei5KZoKeutbzDytHY3+bt9PsCCRAXHiUsMzpBSE9WXWRrcnmAh46VnKOqsbi/xs3U2+Lp8Pf+BQwTGiEoLzY9REtSWWBnbnV8g4qRmJ+mrbS7wsnQ197l7PP6AQgPFh0kKzI5QEdOVVxjanF4f4aNlJuiqbC3vsXM09rh6O/2/QQLEhkgJy41PENKUVhfZm10e4KJkJeepayzusHIz9bd5Ovy+QAHDhUcIyoxOD9GTVRbYmlwd36FjJOaoaivtr3Ey9LZ4Ofu9fwDChEYHyYtNDtCSVBXXmVsc3qBiI+WnaSrsrnAx87V3OPq8fj/Bg==", - "format": "mp4", - }, - }, - ], - }, - ], - "stream": False, - "allowed_openai_params": [], - }, - response=JsonResponse( - content_type="application/json", - body={ - "candidates": [ - { - "content": {"parts": [{"text": "scripted answer 5f5da5957185"}], "role": "model"}, - "finishReason": "STOP", - "index": 0, - } - ], - "usageMetadata": { - "promptTokenCount": 8060, - "candidatesTokenCount": 300, - "totalTokenCount": 8360, - "promptTokensDetails": [ - {"modality": "TEXT", "tokenCount": 140}, - {"modality": "VIDEO", "tokenCount": 7920}, - ], - }, - "modelVersion": "gemini-3.8-flash", - }, - ), - expected=ExactExpected( - spend=0.0059192, input_cost=0.0049832, output_cost=0.000936, prompt_tokens=8060, completion_tokens=300 - ), -) GEMINI_3_8_FLASH_SERVICE_TIER_FLEX: Final = CostTrackingTestCase( - name="gemini-3.8-flash-service_tier_flex", + name="vertex_ai/gemini-3.8-flash-service_tier_flex", covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="gemini-3.8-flash", + model="vertex_ai/gemini-3.8-flash", request={ "model": "$MODEL", "messages": [ @@ -2014,9 +1837,9 @@ GEMINI_3_8_FLASH_SERVICE_TIER_FLEX: Final = CostTrackingTestCase( ) GEMINI_3_8_FLASH_SERVICE_TIER_PRIORITY: Final = CostTrackingTestCase( - name="gemini-3.8-flash-service_tier_priority", + name="vertex_ai/gemini-3.8-flash-service_tier_priority", covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="gemini-3.8-flash", + model="vertex_ai/gemini-3.8-flash", request={ "model": "$MODEL", "messages": [ @@ -2068,9 +1891,9 @@ GEMINI_3_8_FLASH_SERVICE_TIER_PRIORITY: Final = CostTrackingTestCase( ) GEMINI_3_8_FLASH_WEB_SEARCH_PER_PROMPT: Final = CostTrackingTestCase( - name="gemini-3.8-flash-web_search_per_prompt", + name="vertex_ai/gemini-3.8-flash-web_search_per_prompt", covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="gemini-3.8-flash", + model="vertex_ai/gemini-3.8-flash", request={ "model": "$MODEL", "messages": [ @@ -2118,14 +1941,14 @@ GEMINI_3_8_FLASH_WEB_SEARCH_PER_PROMPT: Final = CostTrackingTestCase( }, ), expected=ExactExpected( - spend=0.03724224, input_cost=0.0009568, output_cost=0.00128544, prompt_tokens=1840, completion_tokens=412 + spend=0.10724224, input_cost=0.0009568, output_cost=0.00128544, prompt_tokens=1840, completion_tokens=412 ), ) GEMINI_3_8_FLASH_GOOGLE_MAPS_GROUNDING: Final = CostTrackingTestCase( - name="gemini-3.8-flash-google_maps_grounding", + name="vertex_ai/gemini-3.8-flash-google_maps_grounding", covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="gemini-3.8-flash", + model="vertex_ai/gemini-3.8-flash", request={ "model": "$MODEL", "messages": [ @@ -2182,9 +2005,9 @@ GEMINI_3_8_FLASH_GOOGLE_MAPS_GROUNDING: Final = CostTrackingTestCase( ) GEMINI_3_8_FLASH_FALLBACK_IMAGE_TOKENS_AT_INPUT_RATE: Final = CostTrackingTestCase( - name="gemini-3.8-flash-fallback_image_tokens_at_input_rate", + name="vertex_ai/gemini-3.8-flash-fallback_image_tokens_at_input_rate", covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="gemini-3.8-flash", + model="vertex_ai/gemini-3.8-flash", request={ "model": "$MODEL", "messages": [ @@ -2245,9 +2068,9 @@ GEMINI_3_8_FLASH_FALLBACK_IMAGE_TOKENS_AT_INPUT_RATE: Final = CostTrackingTestCa ) GEMINI_3_8_FLASH_STREAM: Final = CostTrackingTestCase( - name="gemini-3.8-flash-stream", + name="vertex_ai/gemini-3.8-flash-stream", covers="quota_management.spend_tracking.scripted_wire.logs_cost", - model="gemini-3.8-flash", + model="vertex_ai/gemini-3.8-flash", request={ "model": "$MODEL", "messages": [ @@ -2305,9 +2128,9 @@ GEMINI_3_8_FLASH_STREAM: Final = CostTrackingTestCase( ) GEMINI_3_8_FLASH_STREAM_NO_USAGE: Final = CostTrackingTestCase( - name="gemini-3.8-flash-stream_no_usage", + name="vertex_ai/gemini-3.8-flash-stream_no_usage", covers="quota_management.spend_tracking.scripted_wire.logs_cost", - model="gemini-3.8-flash", + model="vertex_ai/gemini-3.8-flash", request={ "model": "$MODEL", "messages": [ @@ -2357,9 +2180,9 @@ GEMINI_3_8_FLASH_STREAM_NO_USAGE: Final = CostTrackingTestCase( ) GEMINI_3_8_FLASH_STREAM_NO_USAGE_TOOL_CALL: Final = CostTrackingTestCase( - name="gemini-3.8-flash-stream_no_usage_tool_call", + name="vertex_ai/gemini-3.8-flash-stream_no_usage_tool_call", covers="quota_management.spend_tracking.scripted_wire.logs_cost", - model="gemini-3.8-flash", + model="vertex_ai/gemini-3.8-flash", request={ "model": "$MODEL", "messages": [ @@ -2441,9 +2264,9 @@ GEMINI_3_8_FLASH_STREAM_NO_USAGE_TOOL_CALL: Final = CostTrackingTestCase( ) GEMINI_3_8_FLASH_STREAM_NO_USAGE_IMAGE_INPUT: Final = CostTrackingTestCase( - name="gemini-3.8-flash-stream_no_usage_image_input", + name="vertex_ai/gemini-3.8-flash-stream_no_usage_image_input", covers="quota_management.spend_tracking.scripted_wire.logs_cost", - model="gemini-3.8-flash", + model="vertex_ai/gemini-3.8-flash", request={ "model": "$MODEL", "messages": [ @@ -2500,9 +2323,9 @@ GEMINI_3_8_FLASH_STREAM_NO_USAGE_IMAGE_INPUT: Final = CostTrackingTestCase( ) GEMINI_3_8_FLASH_PROMPT_BLOCKED: Final = CostTrackingTestCase( - name="gemini-3.8-flash-prompt_blocked", + name="vertex_ai/gemini-3.8-flash-prompt_blocked", covers="quota_management.spend_tracking.scripted_wire.logs_cost", - model="gemini-3.8-flash", + model="vertex_ai/gemini-3.8-flash", request={ "model": "$MODEL", "messages": [ @@ -2550,9 +2373,9 @@ GEMINI_3_8_FLASH_PROMPT_BLOCKED: Final = CostTrackingTestCase( ) GEMINI_3_8_FLASH_STREAM_PROMPT_BLOCKED: Final = CostTrackingTestCase( - name="gemini-3.8-flash-stream_prompt_blocked", + name="vertex_ai/gemini-3.8-flash-stream_prompt_blocked", covers="quota_management.spend_tracking.scripted_wire.logs_cost", - model="gemini-3.8-flash", + model="vertex_ai/gemini-3.8-flash", request={ "model": "$MODEL", "messages": [ @@ -2607,9 +2430,9 @@ GEMINI_3_8_FLASH_STREAM_PROMPT_BLOCKED: Final = CostTrackingTestCase( ) GEMINI_3_8_FLASH_RESPONSE_MODEL_OVERRIDE: Final = CostTrackingTestCase( - name="gemini-3.8-flash-response_model_override", + name="vertex_ai/gemini-3.8-flash-response_model_override", covers="quota_management.spend_tracking.scripted_wire.logs_cost", - model="gemini-3.8-flash", + model="vertex_ai/gemini-3.8-flash", request={ "model": "$MODEL", "messages": [ @@ -2660,9 +2483,9 @@ GEMINI_3_8_FLASH_RESPONSE_MODEL_OVERRIDE: Final = CostTrackingTestCase( ) GEMINI_3_8_FLASH_STREAM_RESPONSE_MODEL_OVERRIDE: Final = CostTrackingTestCase( - name="gemini-3.8-flash-stream_response_model_override", + name="vertex_ai/gemini-3.8-flash-stream_response_model_override", covers="quota_management.spend_tracking.scripted_wire.logs_cost", - model="gemini-3.8-flash", + model="vertex_ai/gemini-3.8-flash", request={ "model": "$MODEL", "messages": [ @@ -2720,9 +2543,9 @@ GEMINI_3_8_FLASH_STREAM_RESPONSE_MODEL_OVERRIDE: Final = CostTrackingTestCase( ) GEMINI_3_8_FLASH_TOOL_CALL: Final = CostTrackingTestCase( - name="gemini-3.8-flash-tool_call", + name="vertex_ai/gemini-3.8-flash-tool_call", covers="quota_management.spend_tracking.scripted_wire.logs_cost", - model="gemini-3.8-flash", + model="vertex_ai/gemini-3.8-flash", request={ "model": "$MODEL", "messages": [ @@ -2807,9 +2630,9 @@ GEMINI_3_8_FLASH_TOOL_CALL: Final = CostTrackingTestCase( ) GEMINI_3_8_FLASH_STREAM_TOOL_CALL: Final = CostTrackingTestCase( - name="gemini-3.8-flash-stream_tool_call", + name="vertex_ai/gemini-3.8-flash-stream_tool_call", covers="quota_management.spend_tracking.scripted_wire.logs_cost", - model="gemini-3.8-flash", + model="vertex_ai/gemini-3.8-flash", request={ "model": "$MODEL", "messages": [ @@ -2901,9 +2724,9 @@ GEMINI_3_8_FLASH_STREAM_TOOL_CALL: Final = CostTrackingTestCase( ) GEMINI_3_8_FLASH_STREAM_FULL_USAGE: Final = CostTrackingTestCase( - name="gemini-3.8-flash-stream_full_usage", + name="vertex_ai/gemini-3.8-flash-stream_full_usage", covers="quota_management.spend_tracking.scripted_wire.logs_cost", - model="gemini-3.8-flash", + model="vertex_ai/gemini-3.8-flash", request={ "model": "$MODEL", "messages": [ @@ -2965,7 +2788,7 @@ GEMINI_3_8_FLASH_STREAM_FULL_USAGE: Final = CostTrackingTestCase( ), ), expected=ExactExpected( - spend=0.006992128, input_cost=0.001619488, output_cost=0.00537264, prompt_tokens=8314, completion_tokens=1592 + spend=0.005775328, input_cost=0.001276288, output_cost=0.00449904, prompt_tokens=8314, completion_tokens=1592 ), ) @@ -3301,26 +3124,182 @@ VERTEX_AI_CLAUDE_FABLE_5_1_TOOL_CALL: Final = _vertex_claude_shape_case( GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT: Final = _gemini_case( - GEMINI_3_8_FLASH_INPUT_TEXT, "gemini-3.1-pro-preview-input_text", "gemini-3.1-pro-preview", + GEMINI_3_8_FLASH_INPUT_TEXT, "vertex_ai/gemini-3.1-pro-preview-input_text", "vertex_ai/gemini-3.1-pro-preview", ExactExpected(spend=0.0094864, input_cost=0.004048, output_cost=0.0054384, prompt_tokens=1840, completion_tokens=412) ) GEMINI_3_1_PRO_PREVIEW_STREAM: Final = _gemini_case( - GEMINI_3_8_FLASH_STREAM, "gemini-3.1-pro-preview-stream", "gemini-3.1-pro-preview", + GEMINI_3_8_FLASH_STREAM, "vertex_ai/gemini-3.1-pro-preview-stream", "vertex_ai/gemini-3.1-pro-preview", ExactExpected(spend=0.0094864, input_cost=0.004048, output_cost=0.0054384, prompt_tokens=1840, completion_tokens=412) ) GEMINI_3_1_PRO_PREVIEW_CACHE_READ: Final = _gemini_case( - GEMINI_3_8_FLASH_CACHE_READ, "gemini-3.1-pro-preview-cache_read", "gemini-3.1-pro-preview", + GEMINI_3_8_FLASH_CACHE_READ, "vertex_ai/gemini-3.1-pro-preview-cache_read", "vertex_ai/gemini-3.1-pro-preview", ExactExpected(spend=0.00912736, input_cost=0.00411136, output_cost=0.005016, prompt_tokens=12928, completion_tokens=380) ) GEMINI_3_1_PRO_PREVIEW_REASONING: Final = _gemini_case( - GEMINI_3_1_PRO_REASONING, "gemini-3.1-pro-preview-reasoning", "gemini-3.1-pro-preview", + GEMINI_3_1_PRO_REASONING, "vertex_ai/gemini-3.1-pro-preview-reasoning", "vertex_ai/gemini-3.1-pro-preview", ExactExpected(spend=0.056056, input_cost=0.002728, output_cost=0.053328, prompt_tokens=1240, completion_tokens=4040) ) GEMINI_3_1_PRO_PREVIEW_TOOL_CALL: Final = _gemini_case( - GEMINI_3_8_FLASH_TOOL_CALL, "gemini-3.1-pro-preview-tool_call", "gemini-3.1-pro-preview", + GEMINI_3_8_FLASH_TOOL_CALL, "vertex_ai/gemini-3.1-pro-preview-tool_call", "vertex_ai/gemini-3.1-pro-preview", ExactExpected(spend=0.0094864, input_cost=0.004048, output_cost=0.0054384, prompt_tokens=1840, completion_tokens=412) ) +GEMINI_4_ARGON_INPUT_TEXT: Final = GEMINI_3_8_FLASH_INPUT_TEXT.model_copy( + update={ + "name": "vertex_ai/gemini-4-argon-input_text", + "model": "vertex_ai/gemini-4-argon", + "response": JsonResponse( + content_type="application/json", + body={ + "candidates": [ + { + "content": {"parts": [{"text": "Hello."}], "role": "model"}, + "finishReason": "STOP", + "index": 0, + } + ], + "usageMetadata": { + "promptTokenCount": 37, + "candidatesTokenCount": 6, + "totalTokenCount": 43, + "trafficType": "ON_DEMAND", + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 37}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 6}], + }, + "modelVersion": "gemini-4-argon", + }, + ), + "expected": ExactExpected( + spend=0.0001608, input_cost=0.0000888, output_cost=0.000072, prompt_tokens=37, completion_tokens=6 + ), + } +) +GEMINI_4_ARGON_STREAM: Final = GEMINI_3_8_FLASH_STREAM.model_copy( + update={ + "name": "vertex_ai/gemini-4-argon-stream", + "model": "vertex_ai/gemini-4-argon", + "response": SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "candidates": [{"content": {"role": "model", "parts": [{"text": "Hello."}]}}], + "usageMetadata": {"trafficType": "ON_DEMAND"}, + "modelVersion": "gemini-4-argon", + }, + { + "candidates": [{"content": {"role": "model", "parts": [{"text": ""}]}, "finishReason": "STOP"}], + "usageMetadata": { + "promptTokenCount": 37, + "candidatesTokenCount": 6, + "totalTokenCount": 43, + "trafficType": "ON_DEMAND", + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 37}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 6}], + }, + "modelVersion": "gemini-4-argon", + }, + ), + ), + "expected": ExactExpected( + spend=0.0001608, input_cost=0.0000888, output_cost=0.000072, prompt_tokens=37, completion_tokens=6 + ), + } +) +GEMINI_4_ARGON_CACHE_READ: Final = GEMINI_3_8_FLASH_CACHE_READ.model_copy( + update={ + "name": "vertex_ai/gemini-4-argon-cache_read", + "model": "vertex_ai/gemini-4-argon", + "response": JsonResponse( + content_type="application/json", + body={ + "candidates": [ + {"content": {"parts": [{"text": "Understood."}], "role": "model"}, "finishReason": "STOP"} + ], + "usageMetadata": { + "promptTokenCount": 13532, + "candidatesTokenCount": 21, + "totalTokenCount": 13553, + "cachedContentTokenCount": 11478, + "trafficType": "ON_DEMAND", + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 13532}], + "cacheTokensDetails": [{"modality": "TEXT", "tokenCount": 11478}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 21}], + }, + "modelVersion": "gemini-4-argon", + }, + ), + "expected": ExactExpected( + spend=0.00655896, input_cost=0.00630696, output_cost=0.000252, prompt_tokens=13532, completion_tokens=21 + ), + } +) +GEMINI_4_ARGON_REASONING: Final = GEMINI_3_1_PRO_REASONING.model_copy( + update={ + "name": "vertex_ai/gemini-4-argon-reasoning", + "model": "vertex_ai/gemini-4-argon", + "response": JsonResponse( + content_type="application/json", + body={ + "candidates": [ + { + "content": { + "parts": [{"text": "Hello.", "thought": True}, {"text": "Hello."}], + "role": "model", + }, + "finishReason": "STOP", + "index": 0, + } + ], + "usageMetadata": { + "promptTokenCount": 37, + "candidatesTokenCount": 6, + "totalTokenCount": 162, + "trafficType": "ON_DEMAND", + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 37}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 6}], + "thoughtsTokenCount": 119, + }, + "modelVersion": "gemini-4-argon", + }, + ), + "expected": ExactExpected( + spend=0.0015888, input_cost=0.0000888, output_cost=0.0015, prompt_tokens=37, completion_tokens=125 + ), + } +) +GEMINI_4_ARGON_TOOL_CALL: Final = GEMINI_3_8_FLASH_TOOL_CALL.model_copy( + update={ + "name": "vertex_ai/gemini-4-argon-tool_call", + "model": "vertex_ai/gemini-4-argon", + "response": JsonResponse( + content_type="application/json", + body={ + "candidates": [ + { + "content": { + "parts": [{"functionCall": {"name": "report_weather", "args": {"city": "Paris"}}}], + "role": "model", + }, + "finishReason": "STOP", + "index": 0, + } + ], + "usageMetadata": { + "promptTokenCount": 90, + "candidatesTokenCount": 22, + "totalTokenCount": 112, + "trafficType": "ON_DEMAND", + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 90}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 22}], + }, + "modelVersion": "gemini-4-argon", + }, + ), + "expected": ExactExpected( + spend=0.00048, input_cost=0.000216, output_cost=0.000264, prompt_tokens=90, completion_tokens=22 + ), + } +) CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GEMINI_3_1_PRO_INPUT_TEXT, @@ -3348,9 +3327,6 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GEMINI_3_1_PRO_STREAM_FULL_USAGE, GEMINI_3_8_FLASH_INPUT_TEXT, GEMINI_3_8_FLASH_CACHE_READ, - GEMINI_3_8_FLASH_AUDIO_INPUT, - GEMINI_3_8_FLASH_AUDIO_OUTPUT, - GEMINI_3_8_FLASH_VIDEO_INPUT, GEMINI_3_8_FLASH_SERVICE_TIER_FLEX, GEMINI_3_8_FLASH_SERVICE_TIER_PRIORITY, GEMINI_3_8_FLASH_WEB_SEARCH_PER_PROMPT, @@ -3396,6 +3372,11 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GEMINI_3_1_PRO_PREVIEW_CACHE_READ, GEMINI_3_1_PRO_PREVIEW_REASONING, GEMINI_3_1_PRO_PREVIEW_TOOL_CALL, + GEMINI_4_ARGON_INPUT_TEXT, + GEMINI_4_ARGON_STREAM, + GEMINI_4_ARGON_CACHE_READ, + GEMINI_4_ARGON_REASONING, + GEMINI_4_ARGON_TOOL_CALL, ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( @@ -3417,6 +3398,7 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( StreamParityTestCase(plain=VERTEX_AI_CLAUDE_OPUS_5_5_INPUT_TEXT, streamed=VERTEX_AI_CLAUDE_OPUS_5_5_STREAM), StreamParityTestCase(plain=VERTEX_AI_CLAUDE_FABLE_5_1_INPUT_TEXT, streamed=VERTEX_AI_CLAUDE_FABLE_5_1_STREAM), StreamParityTestCase(plain=GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT, streamed=GEMINI_3_1_PRO_PREVIEW_STREAM), + StreamParityTestCase(plain=GEMINI_4_ARGON_INPUT_TEXT, streamed=GEMINI_4_ARGON_STREAM), ) BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( @@ -3425,4 +3407,5 @@ BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( ("/v1/responses", "/v1/messages"), ), *bridge_pairs((GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT,), ("/v1/responses",)), + *bridge_pairs((GEMINI_4_ARGON_INPUT_TEXT,), ("/v1/responses",)), ) diff --git a/tests/integration/cost_calculation/catalog.py b/tests/integration/cost_calculation/catalog.py index 66eb7497bd7..568a89ad967 100644 --- a/tests/integration/cost_calculation/catalog.py +++ b/tests/integration/cost_calculation/catalog.py @@ -60,6 +60,7 @@ REALTIME_CASES: Final[tuple[RealtimeCostCase, ...]] = batch_realtime.REALTIME_CA BRIDGE_COST_PARITY_CASES: Final[tuple[BridgeCostParityTestCase, ...]] = ( *openai.BRIDGE_COST_PARITY, *bedrock_converse.BRIDGE_COST_PARITY, + *bedrock_mantle.BRIDGE_COST_PARITY, *azure.BRIDGE_COST_PARITY, *anthropic.BRIDGE_COST_PARITY, *gemini.BRIDGE_COST_PARITY, diff --git a/tests/integration/cost_calculation/cost_map.py b/tests/integration/cost_calculation/cost_map.py index 686636dacbd..5e2b99a113f 100644 --- a/tests/integration/cost_calculation/cost_map.py +++ b/tests/integration/cost_calculation/cost_map.py @@ -492,18 +492,14 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "cache_read_input_token_cost_flex": 2.5e-08, "cache_read_input_token_cost_priority": 9e-08, "google_maps_grounding_cost_per_query": 0.025, - "input_cost_per_audio_token": 1e-06, - "input_cost_per_image_token": 5.5e-07, "input_cost_per_token": 5e-07, "input_cost_per_token_flex": 2.5e-07, "input_cost_per_token_priority": 6.25e-07, - "input_cost_per_video_token": 6e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", - "output_cost_per_audio_token": 6e-06, "output_cost_per_reasoning_token": 3.5e-06, "output_cost_per_token": 3e-06, "output_cost_per_token_flex": 1.5e-06, @@ -536,7 +532,7 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "supports_function_calling": True, "web_search_billing_unit": "per_query", }, - "gemini-3.1-pro-preview": { + "vertex_ai/gemini-3.1-pro-preview": { "cache_creation_input_token_cost_above_200k_tokens": 2.75e-07, "cache_read_input_token_cost": 2.2e-07, "cache_read_input_token_cost_above_200k_tokens": 4.4e-07, @@ -547,7 +543,7 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "input_cost_per_token_above_200k_tokens_priority": 7.92e-06, "input_cost_per_token_flex": 1.1e-06, "input_cost_per_token_priority": 3.96e-06, - "litellm_provider": "vertex_ai-language-models", + "litellm_provider": "vertex_ai", "max_input_tokens": 1048576, "max_output_tokens": 65536, "max_tokens": 65536, @@ -558,33 +554,70 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "output_cost_per_token_above_200k_tokens_priority": 3.564e-05, "output_cost_per_token_flex": 6.6e-06, "output_cost_per_token_priority": 2.376e-05, + "google_maps_grounding_cost_per_query": 0.0154, + "search_context_cost_per_query": { + "search_context_size_high": 0.0154, + "search_context_size_low": 0.0154, + "search_context_size_medium": 0.0154, + }, "supports_function_calling": True, "supports_prompt_caching": True, "supports_reasoning": True, }, - "gemini-3.8-flash": { + "vertex_ai/gemini-3.8-flash": { "cache_read_input_token_cost": 5.2e-08, "cache_read_input_token_cost_flex": 2.6e-08, "cache_read_input_token_cost_priority": 9.36e-08, "google_maps_grounding_cost_per_query": 0.025, - "input_cost_per_audio_token": 1.04e-06, "input_cost_per_token": 5.2e-07, "input_cost_per_token_flex": 2.6e-07, "input_cost_per_token_priority": 6.5e-07, - "input_cost_per_video_token": 6.2e-07, - "litellm_provider": "vertex_ai-language-models", + "litellm_provider": "vertex_ai", "max_input_tokens": 1048576, "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", - "output_cost_per_audio_token": 6.24e-06, "output_cost_per_token": 3.12e-06, "output_cost_per_reasoning_token": 2.6e-06, "output_cost_per_token_flex": 1.56e-06, "output_cost_per_token_priority": 3.9e-06, - "search_context_cost_per_query": {"search_context_size_medium": 0.035}, + "search_context_cost_per_query": { + "search_context_size_high": 0.035, + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035, + }, "supports_function_calling": True, - "web_search_billing_unit": "per_prompt", + "supports_prompt_caching": True, + "supports_reasoning": True, + "web_search_billing_unit": "per_query", + }, + "vertex_ai/gemini-4-argon": { + "cache_read_input_token_cost": 1.2e-07, + "cache_read_input_token_cost_flex": 6e-08, + "cache_read_input_token_cost_priority": 2.16e-07, + "google_maps_grounding_cost_per_query": 0.0168, + "input_cost_per_audio_token": 2.4e-06, + "input_cost_per_token": 2.4e-06, + "input_cost_per_token_flex": 1.2e-06, + "input_cost_per_token_priority": 4.32e-06, + "litellm_provider": "vertex_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_reasoning_token": 1.2e-05, + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_flex": 6e-06, + "output_cost_per_token_priority": 2.16e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.0168, + "search_context_size_low": 0.0168, + "search_context_size_medium": 0.0168, + }, + "supports_function_calling": True, + "supports_prompt_caching": True, + "supports_reasoning": True, + "web_search_billing_unit": "per_query", }, "together_ai/moonshotai/Kimi-K3": { "input_cost_per_token": 1.15e-06, @@ -1078,6 +1111,126 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "supports_reasoning": True, "supports_prompt_caching": True, "supported_endpoints": ["/v1/chat/completions", "/v1/responses"], + "use_openai_responses_path": True, + }, + "bedrock_mantle/openai.gpt-6.1-sol": { + "litellm_provider": "bedrock_mantle", + "mode": "responses", + "input_cost_per_token": 2.42e-06, + "output_cost_per_token": 1.21e-05, + "cache_read_input_token_cost": 1.21e-07, + "cache_creation_input_token_cost": 3.025e-06, + "input_cost_per_token_above_272k_tokens": 4.84e-06, + "output_cost_per_token_above_272k_tokens": 1.815e-05, + "cache_read_input_token_cost_above_272k_tokens": 2.42e-07, + "cache_creation_input_token_cost_above_272k_tokens": 6.05e-06, + "input_cost_per_token_ultrafast": 1.452e-05, + "output_cost_per_token_ultrafast": 7.26e-05, + "cache_read_input_token_cost_ultrafast": 7.26e-07, + "cache_creation_input_token_cost_ultrafast": 1.815e-05, + "input_cost_per_token_above_272k_tokens_ultrafast": 2.904e-05, + "output_cost_per_token_above_272k_tokens_ultrafast": 1.089e-04, + "cache_read_input_token_cost_above_272k_tokens_ultrafast": 1.452e-06, + "cache_creation_input_token_cost_above_272k_tokens_ultrafast": 3.63e-05, + "max_tokens": 131072, + "max_input_tokens": 1000000, + "max_output_tokens": 131072, + "supports_function_calling": True, + "supports_none_reasoning_effort": False, + "supports_reasoning": True, + "supports_prompt_caching": True, + "supported_endpoints": ["/v1/chat/completions", "/v1/responses"], + "use_openai_responses_path": True, + }, + "bedrock_mantle/openai.gpt-6-luna": { + "litellm_provider": "bedrock_mantle", + "mode": "responses", + "input_cost_per_token": 1.32e-07, + "output_cost_per_token": 6.6e-07, + "cache_read_input_token_cost": 1.32e-08, + "cache_creation_input_token_cost": 1.65e-07, + "input_cost_per_token_above_272k_tokens": 2.64e-07, + "output_cost_per_token_above_272k_tokens": 9.9e-07, + "cache_read_input_token_cost_above_272k_tokens": 2.64e-08, + "cache_creation_input_token_cost_above_272k_tokens": 3.3e-07, + "max_tokens": 128000, + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "supports_function_calling": True, + "supports_none_reasoning_effort": False, + "supports_reasoning": True, + "supports_prompt_caching": True, + "supported_endpoints": ["/v1/chat/completions", "/v1/responses"], + "use_openai_responses_path": True, + }, + "bedrock_mantle/openai.gpt-6-astra": { + "litellm_provider": "bedrock_mantle", + "mode": "responses", + "input_cost_per_token": 1.43e-05, + "output_cost_per_token": 7.15e-05, + "cache_read_input_token_cost": 1.43e-06, + "cache_creation_input_token_cost": 1.7875e-05, + "input_cost_per_token_above_272k_tokens": 2.86e-05, + "output_cost_per_token_above_272k_tokens": 1.0725e-04, + "cache_read_input_token_cost_above_272k_tokens": 2.86e-06, + "cache_creation_input_token_cost_above_272k_tokens": 3.575e-05, + "input_cost_per_token_ultrafast": 8.58e-05, + "output_cost_per_token_ultrafast": 4.29e-04, + "cache_read_input_token_cost_ultrafast": 8.58e-06, + "cache_creation_input_token_cost_ultrafast": 1.0725e-04, + "input_cost_per_token_above_272k_tokens_ultrafast": 1.716e-04, + "output_cost_per_token_above_272k_tokens_ultrafast": 6.435e-04, + "cache_read_input_token_cost_above_272k_tokens_ultrafast": 1.716e-05, + "cache_creation_input_token_cost_above_272k_tokens_ultrafast": 2.145e-04, + "max_tokens": 128000, + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "supports_function_calling": True, + "supports_none_reasoning_effort": False, + "supports_reasoning": True, + "supports_prompt_caching": True, + "supported_endpoints": ["/v1/chat/completions", "/v1/responses"], + "use_openai_responses_path": True, + }, + "bedrock_mantle/openai.gpt-5.6-terra": { + "litellm_provider": "bedrock_mantle", + "mode": "responses", + "input_cost_per_token": 3.08e-06, + "output_cost_per_token": 1.848e-05, + "cache_read_input_token_cost": 3.08e-07, + "cache_creation_input_token_cost": 3.85e-06, + "input_cost_per_token_above_272k_tokens": 6.16e-06, + "output_cost_per_token_above_272k_tokens": 2.772e-05, + "cache_read_input_token_cost_above_272k_tokens": 6.16e-07, + "cache_creation_input_token_cost_above_272k_tokens": 7.7e-06, + "max_tokens": 128000, + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "supports_function_calling": True, + "supports_reasoning": True, + "supports_prompt_caching": True, + "supported_endpoints": ["/v1/chat/completions", "/v1/responses"], + "use_openai_responses_path": True, + "search_context_cost_per_query": { + "search_context_size_high": 0.0168, + "search_context_size_low": 0.0168, + "search_context_size_medium": 0.0168, + }, + }, + "bedrock_mantle/anthropic.claude-opus-5-5": { + "litellm_provider": "bedrock_mantle", + "mode": "chat", + "input_cost_per_token": 6.6e-06, + "output_cost_per_token": 3.3e-05, + "cache_read_input_token_cost": 3.3e-07, + "cache_creation_input_token_cost": 8.25e-06, + "cache_creation_input_token_cost_above_1hr": 1.32e-05, + "max_tokens": 128000, + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "supports_function_calling": True, + "supports_reasoning": True, + "supports_prompt_caching": True, }, "us.anthropic.claude-opus-5-5": { "cache_creation_input_token_cost_above_1hr": 8.8e-06, @@ -1410,24 +1563,16 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "cache_creation_input_token_cost": 1.625e-07, "cache_creation_input_token_cost_above_272k_tokens": 3.25e-07, "input_cost_per_token": 1.3e-07, - "input_cost_per_token_flex": 6.5e-08, - "input_cost_per_token_priority": 2.6e-07, "litellm_provider": "azure", "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 6.5e-07, - "output_cost_per_token_flex": 3.25e-07, - "output_cost_per_token_priority": 1.3e-06, "supports_function_calling": True, "supports_none_reasoning_effort": True, "input_cost_per_token_above_272k_tokens": 2.6e-07, - "input_cost_per_token_above_272k_tokens_flex": 1.3e-07, - "input_cost_per_token_above_272k_tokens_priority": 5.2e-07, "output_cost_per_token_above_272k_tokens": 1.3e-06, - "output_cost_per_token_above_272k_tokens_flex": 6.5e-07, - "output_cost_per_token_above_272k_tokens_priority": 2.6e-06, "cache_read_input_token_cost_above_272k_tokens": 2.6e-08, }, "azure/gpt-6-astra": { @@ -1435,24 +1580,16 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "cache_creation_input_token_cost": 1.5625e-05, "cache_creation_input_token_cost_above_272k_tokens": 3.125e-05, "input_cost_per_token": 1.25e-05, - "input_cost_per_token_flex": 6.25e-06, - "input_cost_per_token_priority": 2.5e-05, "litellm_provider": "azure", "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 6.25e-05, - "output_cost_per_token_flex": 3.125e-05, - "output_cost_per_token_priority": 0.000125, "supports_function_calling": True, "supports_none_reasoning_effort": True, "input_cost_per_token_above_272k_tokens": 2.5e-05, - "input_cost_per_token_above_272k_tokens_flex": 1.25e-05, - "input_cost_per_token_above_272k_tokens_priority": 5e-05, "output_cost_per_token_above_272k_tokens": 0.000125, - "output_cost_per_token_above_272k_tokens_flex": 6.25e-05, - "output_cost_per_token_above_272k_tokens_priority": 0.00025, "cache_read_input_token_cost_above_272k_tokens": 2.5e-06, }, "azure/gpt-6.1-sol": { @@ -1460,24 +1597,16 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "cache_creation_input_token_cost": 3e-06, "cache_creation_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 2.4e-06, - "input_cost_per_token_flex": 1.2e-06, - "input_cost_per_token_priority": 4.8e-06, "litellm_provider": "azure", "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 1.2e-05, - "output_cost_per_token_flex": 6e-06, - "output_cost_per_token_priority": 2.4e-05, "supports_function_calling": True, "supports_none_reasoning_effort": True, "input_cost_per_token_above_272k_tokens": 4.8e-06, - "input_cost_per_token_above_272k_tokens_flex": 2.4e-06, - "input_cost_per_token_above_272k_tokens_priority": 9.6e-06, "output_cost_per_token_above_272k_tokens": 2.4e-05, - "output_cost_per_token_above_272k_tokens_flex": 1.2e-05, - "output_cost_per_token_above_272k_tokens_priority": 4.8e-05, "cache_read_input_token_cost_above_272k_tokens": 4.8e-07, }, } diff --git a/tests/integration/cost_calculation/cost_tracking_case.py b/tests/integration/cost_calculation/cost_tracking_case.py index 208d95e6a14..3efa97407f7 100644 --- a/tests/integration/cost_calculation/cost_tracking_case.py +++ b/tests/integration/cost_calculation/cost_tracking_case.py @@ -48,6 +48,7 @@ class CostMapEntry(BaseModel): supports_reasoning: bool | None = None supports_prompt_caching: bool | None = None supported_endpoints: tuple[str, ...] | None = None + use_openai_responses_path: bool | None = None supports_bedrock_runtime_chat_completions_tools_with_reasoning: bool | None = None supports_bedrock_runtime_chat_completions_response_format: bool | None = None cache_creation_input_token_cost_above_100k_tokens: float | None = None @@ -572,6 +573,7 @@ _PROVIDER_PREFIXES: Final[Mapping[str, str]] = MappingProxyType( "deepgram": "deepgram", "text-completion-openai": "text-completion-openai", "cohere": "cohere", + "vertex_ai": "vertex_ai", "vertex_ai-language-models": "vertex_ai", "vertex_ai-image-models": "vertex_ai", "vertex_ai-anthropic_models": "vertex_ai", @@ -621,6 +623,9 @@ _LITELLM_PARAMS: Final[Mapping[str, Mapping[str, str]]] = MappingProxyType( "deepgram": MappingProxyType({}), "text-completion-openai": MappingProxyType({}), "cohere": MappingProxyType({}), + "vertex_ai": MappingProxyType( + {"vertex_project": "cc-scripted-project", "vertex_location": "us-central1"} + ), "vertex_ai-language-models": MappingProxyType( {"vertex_project": "cc-scripted-project", "vertex_location": "us-central1"} ), From d6b54ae2b028653efbc8f93feaf71660b4468501 Mon Sep 17 00:00:00 2001 From: Irvallen Saragih Date: Sat, 17 Oct 2026 05:27:47 +0900 Subject: [PATCH 21/58] fix(openai)!: route gpt-transcribe to GPT audio transcription config (#44021) * fix(openai): route gpt-transcribe to GPT audio transcription config * test(openai): move gpt-transcribe tests under tests/unit * fix(openai): bill gpt-transcribe from provider usage seconds when duration is unknown Fall back to the provider-reported usage.seconds when neither a measured file duration nor a response duration is available, so json responses for undecodable uploads are no longer billed at zero. Move the routing test next to the other provider config tests and replace the structure test with a behavior test. * fix(openai): bill provider usage seconds only for gpt transcription models * fix(cost): decide provider-billed transcription seconds from the deployment model, not its pricing id * test(integration): cover gpt-transcribe default request, provider-billed spend, and health probe --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/cost_calculator.py | 18 ++ litellm/utils.py | 2 +- .../test_audio_transcriptions_wire.py | 174 +++++++++++++++++- .../management/test_model_health_check.py | 37 ++++ tests/unit/test_cost_calculator.py | 72 ++++++++ tests/unit/test_utils.py | 27 +++ 6 files changed, 328 insertions(+), 2 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 0e315e97432..05dc2900349 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -86,6 +86,9 @@ from litellm.llms.openai.cost_calculation import ( from litellm.llms.openai.cost_calculation import ( video_output_cost_per_second, ) +from litellm.llms.openai.transcriptions.gpt_transformation import ( + OpenAIGPTAudioTranscriptionConfig, +) from litellm.llms.perplexity.cost_calculator import ( cost_per_token as perplexity_cost_per_token, ) @@ -325,6 +328,15 @@ def _transcription_usage_has_token_details( return (prompt_tokens_val > 0) or (completion_tokens_val > 0) +def _bills_provider_transcription_seconds(model: str | None, custom_llm_provider: str | None) -> bool: + if model is None or custom_llm_provider != LlmProviders.OPENAI.value: + return False + transcription_config: Final = ProviderConfigManager.get_provider_audio_transcription_config( + model=model.removeprefix(f"{LlmProviders.OPENAI.value}/"), provider=LlmProviders.OPENAI + ) + return isinstance(transcription_config, OpenAIGPTAudioTranscriptionConfig) + + OCRPricingField = Literal["ocr_cost_per_page", "ocr_cost_per_credit", "annotation_cost_per_page"] @@ -1739,6 +1751,12 @@ def completion_cost( audio_transcription_file_duration = _hidden.get( "audio_transcription_duration", getattr(completion_response, "duration", 0.0), + ) or ( + completion_response.usage.seconds + if isinstance(completion_response, TranscriptionResponse) + and isinstance(completion_response.usage, TranscriptionUsageDurationObject) + and _bills_provider_transcription_seconds(requested_model, custom_llm_provider) + else 0.0 ) elif call_type in _RERANK_CALL_TYPES: if completion_response is not None and isinstance(completion_response, RerankResponse): diff --git a/litellm/utils.py b/litellm/utils.py index 8aa6eb23c68..5a933d8f2b8 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -9252,7 +9252,7 @@ class ProviderConfigManager: return XAIAudioTranscriptionConfig() elif litellm.LlmProviders.OPENAI == provider: - if "gpt-4o" in model: + if "gpt-4o" in model or (model.startswith("gpt-") and "transcribe" in model): return litellm.OpenAIGPTAudioTranscriptionConfig() else: return litellm.OpenAIWhisperAudioTranscriptionConfig() diff --git a/tests/integration/compatibility/test_audio_transcriptions_wire.py b/tests/integration/compatibility/test_audio_transcriptions_wire.py index 7845e3d2423..02a09bce5fd 100644 --- a/tests/integration/compatibility/test_audio_transcriptions_wire.py +++ b/tests/integration/compatibility/test_audio_transcriptions_wire.py @@ -1,15 +1,18 @@ +import io import json import math import mimetypes import uuid +import wave from email.message import Message from email.parser import BytesParser from email.policy import HTTP +from hashlib import sha256 from typing import Final, cast import httpx import pytest -from integration._support.client import Gateway, eventually +from integration._support.client import JSON_OBJECT, Gateway, eventually, object_value from integration._support.database import read_rows from integration._support.wire import Reply, Request, wire_server from openai import OpenAI @@ -671,3 +674,172 @@ def test_audio_transcription_bills_audio_tokens_at_the_input_rate_when_no_audio_ rel_tol=1e-9, ), response.text assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/v1/audio/transcriptions")] + + +_GPT_TRANSCRIBE_VERBOSE_JSON_REFUSAL: Final = { + "error": { + "message": "response_format 'verbose_json' is not compatible with model 'gpt-transcribe'", + "type": "invalid_request_error", + "param": "response_format", + "code": "unsupported_value", + } +} +_UNDECODABLE_M4A_BYTES: Final = b"\x00\x00\x00\x20ftypM4A \x00\x00\x02\x00M4A isommp42\x00\x00\x00\x08free" +_DEPLOYMENT_COST_PER_SECOND: Final = 0.0003 + + +def _openai_gpt_transcribe_reply(request: Request, reported_seconds: float) -> Reply: + fields: Final = dict(_text_parts(_multipart_parts(request))) + if fields.get("response_format") == "verbose_json": + return Reply(status=400, body=json.dumps(_GPT_TRANSCRIBE_VERBOSE_JSON_REFUSAL).encode()) + return Reply( + body=json.dumps({"text": "hello world", "usage": {"type": "duration", "seconds": reported_seconds}}).encode() + ) + + +def _uploaded_files(request: Request) -> tuple[tuple[str, str, bytes], ...]: + return tuple((name, filename, payload) for name, filename, _, payload in _file_parts(_multipart_parts(request))) + + +def _silent_wav(seconds: float) -> bytes: + output: Final = io.BytesIO() + with wave.open(output, "wb") as wav: + wav.setnchannels(1) + wav.setsampwidth(2) + wav.setframerate(16000) + wav.writeframes(b"\x00\x00" * round(16000 * seconds)) + return output.getvalue() + + +def _advertised_cost_per_second(gateway: Gateway, model: str) -> float: + entries: Final = gateway.get("/v2/model/info", {"model": model})["data"] + assert isinstance(entries, list) and len(entries) == 1, entries + price: Final = object_value(object_value(entries[0])["model_info"])["input_cost_per_second"] + assert isinstance(price, float) and price > 0, entries + return price + + +def _response_cost(response: httpx.Response) -> float: + return float(response.headers["x-litellm-response-cost"]) if "x-litellm-response-cost" in response.headers else 0.0 + + +def _logged_spend(call_id: str) -> float: + rows: Final = eventually( + lambda: read_rows('SELECT spend FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (call_id,)), + lambda values: len(values) == 1, + seconds=70, + ) + return float(str(rows[0]["spend"])) + + +def test_gpt_transcribe_without_response_format_reaches_openai_without_verbose_json(gateway: Gateway) -> None: + upload: Final = _UNDECODABLE_M4A_BYTES + uuid.uuid4().bytes + + with ( + wire_server(lambda request: _openai_gpt_transcribe_reply(request, 2)) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model( + model="openai/gpt-transcribe", api_base=f"{wire.url}/v1", api_key="synthetic-openai-key" + ) + with _sdk(gateway, "/v1", gateway.key) as sdk: + response: Final = sdk.audio.transcriptions.with_raw_response.create( + model=model, file=("a.m4a", upload, "audio/mp4") + ) + requests: Final = wire.drain() + + assert response.status_code == 200, response.text + assert JSON_OBJECT.validate_json(response.content)["text"] == "hello world", response.text + assert [(request.method, request.target) for request in requests] == [("POST", "/v1/audio/transcriptions")] + assert _text_parts(_multipart_parts(requests[0])) == (("model", "gpt-transcribe"),) + assert _uploaded_files(requests[0]) == (("file", "a.m4a", upload),) + + +@pytest.mark.parametrize( + "configured_cost_per_second", + (None, _DEPLOYMENT_COST_PER_SECOND), + ids=("cost_map_price", "deployment_price"), +) +def test_gpt_transcribe_unmeasurable_upload_bills_the_seconds_openai_reports( + gateway: Gateway, configured_cost_per_second: float | None +) -> None: + reported_seconds: Final = 2.0 + upload: Final = _UNDECODABLE_M4A_BYTES + uuid.uuid4().bytes + + with ( + wire_server(lambda request: _openai_gpt_transcribe_reply(request, reported_seconds)) as wire, + gateway.scenario() as scenario, + ): + model: Final = ( + scenario.model( + model="openai/gpt-transcribe", + api_base=f"{wire.url}/v1", + api_key="synthetic-openai-key", + input_cost_per_second=configured_cost_per_second, + ) + if configured_cost_per_second is not None + else scenario.model( + model="openai/gpt-transcribe", api_base=f"{wire.url}/v1", api_key="synthetic-openai-key" + ) + ) + key: Final = scenario.key(models=[model]) + cost_per_second: Final = ( + configured_cost_per_second + if configured_cost_per_second is not None + else _advertised_cost_per_second(gateway, model) + ) + response: Final = gateway.request_multipart( + "/v1/audio/transcriptions", + {"model": model, "response_format": "json"}, + {"file": ("a.m4a", upload, "audio/mp4")}, + key=key, + ) + expected_cost: Final = reported_seconds * cost_per_second + + assert response.status_code == 200, response.text + assert JSON_OBJECT.validate_json(response.content)["text"] == "hello world", response.text + assert math.isclose(_response_cost(response), expected_cost), (_response_cost(response), expected_cost) + assert math.isclose(_logged_spend(response.headers["x-litellm-call-id"]), expected_cost) + key_rows: Final = eventually( + lambda: read_rows( + 'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token = %s', + (sha256(key.encode()).hexdigest(),), + ), + lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) > 0, + seconds=70, + ) + assert math.isclose(float(str(key_rows[0]["spend"])), expected_cost), key_rows + requests: Final = wire.drain() + + assert [(request.method, request.target) for request in requests] == [("POST", "/v1/audio/transcriptions")] + assert _uploaded_files(requests[0]) == (("file", "a.m4a", upload),) + + +def test_gpt_transcribe_measured_upload_bills_the_measured_length_over_reported_seconds(gateway: Gateway) -> None: + measured_seconds: Final = 1.0 + upload: Final = _silent_wav(measured_seconds) + + with ( + wire_server(lambda request: _openai_gpt_transcribe_reply(request, 7)) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model( + model="openai/gpt-transcribe", + api_base=f"{wire.url}/v1", + api_key="synthetic-openai-key", + input_cost_per_second=_DEPLOYMENT_COST_PER_SECOND, + ) + response: Final = gateway.request_multipart( + "/v1/audio/transcriptions", + {"model": model, "response_format": "json"}, + {"file": ("a.wav", upload, _WAV_CONTENT_TYPE)}, + ) + expected_cost: Final = measured_seconds * _DEPLOYMENT_COST_PER_SECOND + + assert response.status_code == 200, response.text + assert math.isclose(_response_cost(response), expected_cost), (_response_cost(response), expected_cost) + assert math.isclose(_logged_spend(response.headers["x-litellm-call-id"]), expected_cost) + requests: Final = wire.drain() + + assert [(request.method, request.target) for request in requests] == [("POST", "/v1/audio/transcriptions")] + assert _uploaded_files(requests[0]) == (("file", "a.wav", upload),) diff --git a/tests/integration/management/test_model_health_check.py b/tests/integration/management/test_model_health_check.py index 2b9ec658d87..9945d16af07 100644 --- a/tests/integration/management/test_model_health_check.py +++ b/tests/integration/management/test_model_health_check.py @@ -1352,3 +1352,40 @@ def test_evaluation_mode_health_check_of_databricks_ai_decide_resolves_the_mode_ }, ) ] + + +def _openai_gpt_transcribe_health_reply(request: Request) -> Reply: + if _multipart_text_fields(_multipart_parts(request)).get("response_format") == "verbose_json": + return Reply( + status=400, + body=json.dumps( + { + "error": { + "message": "response_format 'verbose_json' is not compatible with model 'gpt-transcribe'", + "type": "invalid_request_error", + "param": "response_format", + "code": "unsupported_value", + } + } + ).encode(), + ) + return Reply(body=json.dumps({"text": "healthy", "usage": {"type": "duration", "seconds": 1}}).encode()) + + +def test_health_check_of_gpt_transcribe_probes_without_verbose_json_and_reports_it_healthy(gateway: Gateway) -> None: + with wire_server(_openai_gpt_transcribe_health_reply) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/gpt-transcribe", + api_base=f"{wire.url}/v1", + api_key="synthetic-openai-key", + model_info={"mode": "audio_transcription"}, + ) + report: Final = _health_report(gateway, model) + requests: Final = wire.drain() + + assert (report["healthy_count"], report["unhealthy_count"]) == (1, 0), report + assert requests, report + assert { + (request.method, request.target, tuple(sorted(_multipart_text_fields(_multipart_parts(request)).items()))) + for request in requests + } == {("POST", "/v1/audio/transcriptions", (("model", "gpt-transcribe"),))} diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index cc2a1d62d52..c61ebb3737c 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -6426,6 +6426,78 @@ def test_whisper_azure(): assert round(cost, 5) == round(expected_cost, 5) +@pytest.mark.parametrize( + ("provider", "requested_model", "response_extras", "expected_seconds"), + [ + ("openai", "openai/gpt-fake-transcribe", {}, 2.0), + ("openai", "openai/gpt-fake-transcribe", {"_audio_transcription_duration": 1.35}, 1.35), + ("openai", "openai/gpt-fake-transcribe", {"duration": 1.34}, 1.34), + ("openai", "openai/whisper-fake", {}, 0.0), + ("azure", "gpt-fake-transcribe", {}, 0.0), + ], + ids=[ + "undecodable_upload_bills_provider_seconds", + "measured_duration_still_wins", + "provider_duration_still_wins", + "whisper_models_never_bill_provider_seconds", + "azure_never_bills_provider_seconds", + ], +) +def test_transcription_cost_falls_back_to_provider_billed_seconds( + monkeypatch: pytest.MonkeyPatch, + provider: str, + requested_model: str, + response_extras: Mapping[str, float], + expected_seconds: float, +) -> None: + model: Final = requested_model.removeprefix(f"{provider}/") + monkeypatch.setitem( + litellm.model_cost, + model, + {"input_cost_per_second": 0.01, "litellm_provider": provider, "mode": "audio_transcription"}, + ) + transcription: Final = convert_to_model_response_object( + response_object={ + "text": "What is the weather in Paris?", + "usage": {"type": "duration", "seconds": 2}, + **response_extras, + }, + model_response_object=TranscriptionResponse(), + hidden_params={"model": model, "custom_llm_provider": provider}, + response_type="audio_transcription", + ) + + cost: Final = completion_cost( + completion_response=transcription, model=requested_model, call_type="atranscription" + ) + + assert cost == pytest.approx(0.01 * expected_seconds) + + +def test_custom_priced_transcription_deployment_bills_provider_seconds(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setitem( + litellm.model_cost, + "priced-deployment", + {"input_cost_per_second": 0.000075, "litellm_provider": "openai", "mode": "audio_transcription"}, + ) + transcription: Final = convert_to_model_response_object( + response_object={"text": "What is the weather in Paris?", "usage": {"type": "duration", "seconds": 2}}, + model_response_object=TranscriptionResponse(), + hidden_params={"model": "gpt-fake-transcribe", "custom_llm_provider": "openai"}, + response_type="audio_transcription", + ) + + cost: Final = completion_cost( + completion_response=transcription, + model="openai/gpt-fake-transcribe", + call_type="atranscription", + custom_pricing=True, + router_model_id="priced-deployment", + ) + + assert cost == pytest.approx(0.000075 * 2) + + def test_gpt_image_2_azure_cost_tracking(): azure_image_generation_response: Final = { "created": 1758585600, diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 98025f3a79a..75cec231f16 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -9585,6 +9585,33 @@ def test_get_provider_audio_transcription_config_resolves_for_every_provider() - assert isinstance(configs[LlmProviders.OPENAI], litellm.OpenAIWhisperAudioTranscriptionConfig) +@pytest.mark.parametrize( + ("model", "optional_params", "expected_response_format"), + [ + ("whisper-1", {}, "verbose_json"), + ("my-transcribe-whisper", {}, "verbose_json"), + ("whisper-1", {"response_format": "srt"}, "srt"), + ("gpt-4o-mini-transcribe", {}, None), + ("gpt-transcribe", {}, None), + ("gpt-transcribe", {"response_format": "text"}, "text"), + ], +) +def test_openai_transcription_defaults_verbose_json_only_for_whisper_models( + model: str, optional_params: dict[str, str], expected_response_format: str | None +) -> None: + config: Final = ProviderConfigManager.get_provider_audio_transcription_config( + model=model, provider=LlmProviders.OPENAI + ) + assert config is not None + + request: Final = config.transform_audio_transcription_request( + model=model, audio_file=b"audio", optional_params=dict(optional_params), litellm_params={} + ) + + assert isinstance(request.data, dict) + assert request.data.get("response_format") == expected_response_format + + def test_get_valid_models_from_provider_cache_invalidation(monkeypatch): """ Test that get_valid_models returns the correct models for a given provider From c7b115b56a310c4a0de8ca38295bde680fe7b2a2 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 16 Oct 2026 13:44:24 -0700 Subject: [PATCH 22/58] feat(ui): add the 1 hour cache write cost to the add model and model edit forms (#46830) * feat(ui): add the 1 hour cache write cost to the add model and model edit forms Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(e2e): cover both cache write tiers in the add model and edit forms Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(ui): hoist the expected pricing payload out of the inline assertion Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(ui): share the per-token conversion between both cache write tiers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(ui): derive the touched pricing field type from the pricing field list Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(ui): collapse the duplicated pricing source keys into one map Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ui): convert add model pricing once for all selected models and describe the 1 hour fallback Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(e2e): type the cache write tier deployment read back and fill the 1 hour field Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(e2e): keep the cache write tier spec on the dashboard prettier config Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/types/utils.py | 2 +- .../tests/modelsPage/cacheWriteTiers.spec.ts | 146 ++++++++++++++++++ .../modelsPage/clearCustomPricing.spec.ts | 2 +- .../test_model_management_endpoints.py | 37 +++++ tests/unit/types/test_router.py | 1 + .../src/components/ModelInfoEditForm.tsx | 33 ++-- .../add_model/advanced_settings.test.tsx | 11 ++ .../add_model/advanced_settings.tsx | 20 +++ .../handle_add_model_submit.test.tsx | 55 +++++++ .../add_model/handle_add_model_submit.tsx | 111 +++++++------ .../src/components/model_info_view.test.tsx | 43 +++++- .../src/components/model_info_view.tsx | 21 ++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 + 13 files changed, 409 insertions(+), 75 deletions(-) create mode 100644 tests/e2e/ui/tests/modelsPage/cacheWriteTiers.spec.ts diff --git a/litellm/types/utils.py b/litellm/types/utils.py index a97e35a9021..054a0a1c034 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3900,6 +3900,7 @@ class MirroredPricingParams(LiteLLMBaseModel): output_cost_per_character: float | None = None cache_read_input_token_cost: float | None = None cache_creation_input_token_cost: float | None = None + cache_creation_input_token_cost_above_1hr: float | None = None tiered_pricing: list[dict[str, Any]] | None = None @@ -3927,7 +3928,6 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): input_cost_per_token_priority: float | None = None input_cost_per_token_balanced: float | None = None input_cost_per_token_ultrafast: float | None = None - cache_creation_input_token_cost_above_1hr: float | None = None cache_creation_input_token_cost_above_100k_tokens: float | None = None cache_creation_input_token_cost_above_1hr_above_100k_tokens: float | None = None cache_creation_input_token_cost_above_200k_tokens: float | None = None diff --git a/tests/e2e/ui/tests/modelsPage/cacheWriteTiers.spec.ts b/tests/e2e/ui/tests/modelsPage/cacheWriteTiers.spec.ts new file mode 100644 index 00000000000..7d0f02209cf --- /dev/null +++ b/tests/e2e/ui/tests/modelsPage/cacheWriteTiers.spec.ts @@ -0,0 +1,146 @@ +import { test, expect, type Page as PlaywrightPage } from "@playwright/test"; +import { ADMIN_STORAGE_PATH } from "../../constants"; +import { Role, users } from "../../fixtures/users"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage } from "../../helpers/navigation"; +import { captureRequestBody, readBack } from "../../helpers/roundTrip"; + +const FIVE_MINUTE_PER_MILLION = 3.75; +const ONE_HOUR_PER_MILLION = 7.5; +const perToken = (perMillion: number): number => perMillion / 1_000_000; +const ONE_HOUR_LABEL = /^Cache Write Cost \(1 hour, per 1M tokens\)/; + +type Deployment = { + model_name?: string; + model_info?: { id?: string }; + litellm_params: { + cache_creation_input_token_cost?: number; + cache_creation_input_token_cost_above_1hr?: number; + }; +}; + +const masterKey = users[Role.ProxyAdmin].password; +const auth = { Authorization: `Bearer ${masterKey}` }; + +async function deleteModel(page: PlaywrightPage, id: string) { + await page.request.post("/model/delete", { headers: auth, data: { id } }); +} + +async function readDeployment(page: PlaywrightPage, modelId: string): Promise { + const body = await readBack<{ data: Deployment[] }>(page, `/model/info?litellm_model_id=${modelId}`); + return body.data[0]; +} + +async function selectProvider(page: PlaywrightPage, providerName: string) { + const providerDropdown = page.getByRole("combobox", { name: "Provider", exact: true }); + await providerDropdown.click(); + await providerDropdown.fill(providerName); + await page.getByRole("option").filter({ hasText: providerName }).first().click(); + await expect(providerDropdown).toHaveValue(providerName); +} + +test.describe("Cache write cost tiers", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + let createdModelId: string | null = null; + + test.afterEach(async ({ page }) => { + if (createdModelId) { + await deleteModel(page, createdModelId); + createdModelId = null; + } + }); + + test("the edit form saves the 1 hour rate next to the 5 minute rate and the detail view shows both", async ({ + page, + }) => { + const modelName = `e2e-cache-tiers-edit-${Date.now()}`; + const created = await page.request.post("/model/new", { + headers: auth, + data: { + model_name: modelName, + litellm_params: { + model: "anthropic/claude-sonnet-4-6", + api_key: "sk-e2e-not-used", + cache_creation_input_token_cost: perToken(FIVE_MINUTE_PER_MILLION), + }, + model_info: {}, + }, + }); + expect(created.ok(), `POST /model/new for ${modelName}`).toBe(true); + createdModelId = (await created.json()).model_info?.id; + expect(createdModelId, "model id from /model/new").toBeTruthy(); + + await page.goto("/ui"); + await page.getByText("Models + Endpoints").click(); + const modelIdCell = page.getByTestId(`model-id-${createdModelId}`); + await expect(modelIdCell).toBeVisible({ timeout: 15_000 }); + await modelIdCell.click(); + await expect(page.getByText("Back to Models").first()).toBeVisible({ timeout: 10_000 }); + + await page.getByRole("button", { name: "Edit Settings" }).click(); + const oneHour = page.getByLabel(ONE_HOUR_LABEL); + await oneHour.waitFor({ timeout: 15_000 }); + await oneHour.fill(String(ONE_HOUR_PER_MILLION)); + + const patchBody = await captureRequestBody( + page, + { method: "PATCH", urlIncludes: `/model/${createdModelId}/update` }, + async () => { + await page.getByRole("button", { name: "Save Changes" }).click(); + }, + ); + expect(patchBody.litellm_params.cache_creation_input_token_cost_above_1hr, "1 hour rate sent per token").toBe( + perToken(ONE_HOUR_PER_MILLION), + ); + await expect(page.getByText("Model settings updated successfully")).toBeVisible({ timeout: 10_000 }); + + const stored = await readDeployment(page, createdModelId!); + expect(stored?.litellm_params?.cache_creation_input_token_cost_above_1hr, "1 hour rate persisted").toBe( + perToken(ONE_HOUR_PER_MILLION), + ); + expect(stored?.litellm_params?.cache_creation_input_token_cost, "5 minute rate untouched").toBe( + perToken(FIVE_MINUTE_PER_MILLION), + ); + + await expect(page.getByText(ONE_HOUR_LABEL)).toBeVisible(); + await expect(page.getByText(ONE_HOUR_PER_MILLION.toFixed(4))).toBeVisible({ timeout: 10_000 }); + await expect(page.getByText(FIVE_MINUTE_PER_MILLION.toFixed(4))).toBeVisible(); + }); + + test("Add Model sends both cache write tiers per token", async ({ page }) => { + const publicName = `e2e-cache-tiers-add-${Date.now()}`; + await navigateToPage(page, Page.Models); + await page.getByRole("tab", { name: "Add Model" }).click(); + await selectProvider(page, "Anthropic"); + + await page.getByRole("combobox", { name: "Select models" }).click(); + await page.getByRole("option", { name: "Custom Model Name (Enter below)" }).click(); + await page.keyboard.press("Escape"); + await page.getByPlaceholder("Enter custom model name").fill(publicName); + await page.locator("#api_key").fill("sk-e2e-not-used"); + + await page.getByText("Advanced Settings").click(); + await page.getByRole("switch", { name: "Custom Pricing" }).click(); + await page.getByLabel(/^Input Cost \(per 1M tokens\)/).fill("3"); + await page.getByLabel(/^Output Cost \(per 1M tokens\)/).fill("15"); + await page.getByLabel(/^Cache Write Cost \(per 1M tokens\)/).fill(String(FIVE_MINUTE_PER_MILLION)); + await page.getByLabel(ONE_HOUR_LABEL).fill(String(ONE_HOUR_PER_MILLION)); + + const body = await captureRequestBody(page, { method: "POST", urlIncludes: "/model/new" }, async () => { + await page.getByTestId("add-model-btn").click(); + }); + await expect(page.getByText("created successfully")).toBeVisible({ timeout: 15_000 }); + const listed = await readBack<{ data: Deployment[] }>(page, "/v2/model/info"); + createdModelId = listed.data.find((row) => row.model_name === publicName)?.model_info?.id ?? null; + + expect(body.model_name).toBe(publicName); + expect(body.litellm_params.cache_creation_input_token_cost, "5 minute rate sent per token").toBe( + perToken(FIVE_MINUTE_PER_MILLION), + ); + expect(body.litellm_params.cache_creation_input_token_cost_above_1hr, "1 hour rate sent per token").toBe( + perToken(ONE_HOUR_PER_MILLION), + ); + expect(createdModelId, "the model landed in the database").toBeTruthy(); + }); +}); diff --git a/tests/e2e/ui/tests/modelsPage/clearCustomPricing.spec.ts b/tests/e2e/ui/tests/modelsPage/clearCustomPricing.spec.ts index c532641b238..34a93c7ce14 100644 --- a/tests/e2e/ui/tests/modelsPage/clearCustomPricing.spec.ts +++ b/tests/e2e/ui/tests/modelsPage/clearCustomPricing.spec.ts @@ -87,7 +87,7 @@ test.describe("Clear custom pricing on a deployment", () => { // Both cache fields share the same placeholder ("Defaults to Input Cost if blank"), // so disambiguate via their labels. const cacheReadCost = page.getByLabel(/Cache Read Cost/); - const cacheWriteCost = page.getByLabel(/Cache Write Cost/); + const cacheWriteCost = page.getByLabel("Cache Write Cost (per 1M tokens)", { exact: true }); await inputCost.waitFor({ timeout: 15_000 }); for (const field of [inputCost, outputCost, cacheReadCost, cacheWriteCost]) { await field.click({ clickCount: 3 }); diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index beffe454b8d..ad75344afba 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -3758,6 +3758,43 @@ class TestUpdateDBModelClearPricing: assert "cache_creation_input_token_cost" not in params assert "cache_creation_input_token_cost" not in info + def test_clear_cache_write_1h_cost_removes_from_both_blobs(self): + """The dashboard clears the 1 hour cache-write tier with an explicit null, which only + takes effect for fields in SPECIAL_MODEL_INFO_PARAMS.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_db_model, + ) + from litellm.types.router import ( + Deployment, + LiteLLM_Params, + ModelInfo, + updateLiteLLMParams, + ) + + db_model = Deployment( + model_name="anthropic/*", + litellm_params=LiteLLM_Params( + model="anthropic/*", + cache_creation_input_token_cost=0.00000375, + cache_creation_input_token_cost_above_1hr=0.000006, + ), + model_info=ModelInfo(id="dep-cache-write-1h-0"), + ) + + result = update_db_model( + db_model=db_model, + updated_patch=updateDeployment( + litellm_params=updateLiteLLMParams(cache_creation_input_token_cost_above_1hr=None) + ), + ) + + params = json.loads(result["litellm_params"]) + info = json.loads(result["model_info"]) + assert "cache_creation_input_token_cost_above_1hr" not in params + assert "cache_creation_input_token_cost_above_1hr" not in info + assert params["cache_creation_input_token_cost"] == 0.00000375 + assert info["cache_creation_input_token_cost"] == 0.00000375 + def test_clear_cache_read_preserves_other_pricing(self): """Clearing cache_read must not touch input/output cost overrides.""" from litellm.proxy.management_endpoints.model_management_endpoints import ( diff --git a/tests/unit/types/test_router.py b/tests/unit/types/test_router.py index 6188bb966ee..8092437fbd3 100644 --- a/tests/unit/types/test_router.py +++ b/tests/unit/types/test_router.py @@ -63,6 +63,7 @@ def test_custom_pricing_params_keeps_every_field_it_had(): "output_cost_per_character", "cache_read_input_token_cost", "cache_creation_input_token_cost", + "cache_creation_input_token_cost_above_1hr", "cost_per_second", "input_cost_per_second", "cache_read_input_token_cost_flex", diff --git a/ui/litellm-dashboard/src/components/ModelInfoEditForm.tsx b/ui/litellm-dashboard/src/components/ModelInfoEditForm.tsx index c7860f35100..125857b9a94 100644 --- a/ui/litellm-dashboard/src/components/ModelInfoEditForm.tsx +++ b/ui/litellm-dashboard/src/components/ModelInfoEditForm.tsx @@ -63,20 +63,21 @@ const PTU_EDIT_FIELDS: PtuEditField[] = [ { name: PTU_END_FIELD, label: "PTU Effective To (UTC)", input: "datetime" }, ]; -export type TouchedPricingField = "input_cost" | "output_cost" | "cache_read_cost" | "cache_write_cost"; - -const PRICING_FIELDS: readonly TouchedPricingField[] = [ +const PRICING_FIELDS = [ "input_cost", "output_cost", "cache_read_cost", "cache_write_cost", + "cache_write_cost_1h", ] as const; +export type TouchedPricingField = (typeof PRICING_FIELDS)[number]; -const COST_SOURCES: Record = { - input_cost: { param: "input_cost_per_token", info: "input_cost_per_token" }, - output_cost: { param: "output_cost_per_token", info: "output_cost_per_token" }, - cache_read_cost: { param: "cache_read_input_token_cost", info: "cache_read_input_token_cost" }, - cache_write_cost: { param: "cache_creation_input_token_cost", info: "cache_creation_input_token_cost" }, +const COST_SOURCES: Record = { + input_cost: "input_cost_per_token", + output_cost: "output_cost_per_token", + cache_read_cost: "cache_read_input_token_cost", + cache_write_cost: "cache_creation_input_token_cost", + cache_write_cost_1h: "cache_creation_input_token_cost_above_1hr", }; export interface ModelEditFormValues { @@ -94,6 +95,7 @@ export interface ModelEditFormValues { output_cost?: string | number | null; cache_read_cost?: string | number | null; cache_write_cost?: string | number | null; + cache_write_cost_1h?: string | number | null; ptu_count?: string | number | null; cost_per_ptu_per_hour?: string | number | null; ptu_effective_from?: Dayjs | null; @@ -131,6 +133,7 @@ const modelEditShape = { output_cost: scalar, cache_read_cost: scalar, cache_write_cost: scalar, + cache_write_cost_1h: scalar, ptu_count: scalar, cost_per_ptu_per_hour: scalar, ptu_effective_from: z.custom().nullish(), @@ -243,6 +246,10 @@ export const toModelEditFormValues = (localModelData: any, isWildcardModel: bool localModelData.litellm_params?.cache_creation_input_token_cost, localModelData.model_info?.cache_creation_input_token_cost, ), + cache_write_cost_1h: perMillionTokens( + localModelData.litellm_params?.cache_creation_input_token_cost_above_1hr, + localModelData.model_info?.cache_creation_input_token_cost_above_1hr, + ), cache_control: localModelData.litellm_params?.cache_control_injection_points ? true : false, cache_control_injection_points: localModelData.litellm_params?.cache_control_injection_points || [], model_access_group: Array.isArray(localModelData.model_info?.access_groups) @@ -271,8 +278,8 @@ export const toModelEditFormValues = (localModelData: any, isWildcardModel: bool }); const displayCost = (localModelData: any, field: TouchedPricingField): string => { - const { param, info } = COST_SOURCES[field]; - const rate = localModelData?.litellm_params?.[param] ?? localModelData?.model_info?.[info]; + const key = COST_SOURCES[field]; + const rate = localModelData?.litellm_params?.[key] ?? localModelData?.model_info?.[key]; return rate != null ? (Number(rate) * 1_000_000).toFixed(4) : "Not Set"; }; @@ -527,6 +534,12 @@ const ModelInfoEditForm: React.FC = ({ "Defaults to Input Cost if blank", "If left blank on save, defaults to Input Cost (backend falls back to input_cost_per_token).", )} + {pricingField( + "cache_write_cost_1h", + "Cache Write Cost (1 hour, per 1M tokens)", + "Defaults to the cost map 1 hour rate, else Cache Write Cost", + "Rate for prompt cache writes with a 1 hour TTL. If left blank on save, known models use their cost map 1 hour rate, otherwise Cache Write Cost.", + )} {textField("api_base", "API Base", "Enter API base", localModelData.litellm_params?.api_base)} {textField( diff --git a/ui/litellm-dashboard/src/components/add_model/advanced_settings.test.tsx b/ui/litellm-dashboard/src/components/add_model/advanced_settings.test.tsx index ceb9d1ca713..15a13219d18 100644 --- a/ui/litellm-dashboard/src/components/add_model/advanced_settings.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/advanced_settings.test.tsx @@ -52,6 +52,17 @@ describe("AdvancedSettings", () => { }); }); + it("offers both cache write tiers once custom pricing is on", async () => { + renderAdvancedSettings(); + act(() => { + fireEvent.click(screen.getByText("Advanced Settings")); + }); + fireEvent.click(await screen.findByRole("switch", { name: "Custom Pricing" })); + + expect(await screen.findByLabelText(/^Cache Write Cost \(per 1M tokens\)/)).toBeInTheDocument(); + expect(screen.getByLabelText(/^Cache Write Cost \(1 hour, per 1M tokens\)/)).toBeInTheDocument(); + }); + it("hides every PTU field when PTU cost attribution is disabled", async () => { renderAdvancedSettings(); act(() => { diff --git a/ui/litellm-dashboard/src/components/add_model/advanced_settings.tsx b/ui/litellm-dashboard/src/components/add_model/advanced_settings.tsx index eb12e49ca89..6025db8e2e2 100644 --- a/ui/litellm-dashboard/src/components/add_model/advanced_settings.tsx +++ b/ui/litellm-dashboard/src/components/add_model/advanced_settings.tsx @@ -49,6 +49,7 @@ const USAGE_COST_FIELDS = [ "output_cost_per_token", "cache_read_input_token_cost", "cache_creation_input_token_cost", + "cache_creation_input_token_cost_above_1hr", "input_cost_per_second", ]; @@ -393,6 +394,25 @@ const AdvancedSettings: React.FC = ({ /> )} + + {(control) => ( + + )} + ) : ( { expect(deployment.modelInfoObj.team_id).toBe("team-123"); }); + it("converts both cache write tiers from per-million to per-token", async () => { + const formValues = { + model_mappings: [{ public_name: "claude-sonnet", litellm_model: "anthropic/claude-sonnet-4-6" }], + model_name: "anthropic/claude-sonnet-4-6", + custom_pricing: true, + input_cost_per_token: "3", + output_cost_per_token: "15", + cache_creation_input_token_cost: "3.75", + cache_creation_input_token_cost_above_1hr: "6", + }; + + const [deployment] = (await prepareModelAddRequest({ ...formValues }, "token", null))!; + + expect(deployment.litellmParamsObj.cache_creation_input_token_cost).toBeCloseTo(0.00000375, 12); + expect(deployment.litellmParamsObj.cache_creation_input_token_cost_above_1hr).toBeCloseTo(0.000006, 12); + }); + + it("converts the pricing once when several models are added together", async () => { + const formValues = { + model_mappings: [ + { public_name: "claude-sonnet", litellm_model: "anthropic/claude-sonnet-4-6" }, + { public_name: "claude-opus", litellm_model: "anthropic/claude-opus-4-2" }, + ], + model_name: "anthropic/claude-sonnet-4-6", + custom_pricing: true, + input_cost_per_token: "3", + cache_creation_input_token_cost: "3.75", + cache_creation_input_token_cost_above_1hr: "6", + }; + + const deployments = (await prepareModelAddRequest({ ...formValues }, "token", null))!; + + expect(deployments).toHaveLength(2); + for (const deployment of deployments) { + expect(deployment.litellmParamsObj.input_cost_per_token).toBeCloseTo(0.000003, 12); + expect(deployment.litellmParamsObj.cache_creation_input_token_cost).toBeCloseTo(0.00000375, 12); + expect(deployment.litellmParamsObj.cache_creation_input_token_cost_above_1hr).toBeCloseTo(0.000006, 12); + } + }); + + it("leaves the 1 hour cache write tier unset when its field is blank", async () => { + const formValues = { + model_mappings: [{ public_name: "claude-sonnet", litellm_model: "anthropic/claude-sonnet-4-6" }], + model_name: "anthropic/claude-sonnet-4-6", + custom_pricing: true, + cache_creation_input_token_cost: "3.75", + cache_creation_input_token_cost_above_1hr: "", + }; + + const [deployment] = (await prepareModelAddRequest({ ...formValues }, "token", null))!; + + expect(deployment.litellmParamsObj).not.toHaveProperty("cache_creation_input_token_cost_above_1hr"); + expect(deployment.litellmParamsObj.cache_creation_input_token_cost).toBeCloseTo(0.00000375, 12); + }); + it("uses a lowercase fallback for unrecognized custom providers", async () => { const fallbackValues = { model_mappings: [ diff --git a/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.tsx b/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.tsx index 5f2747e0415..70b499b108a 100644 --- a/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.tsx +++ b/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.tsx @@ -24,6 +24,65 @@ export const prepareModelAddRequest = async (formValues: Record, ac formValues["model"] = wildcardModel; } + // Handle pricing conversion before processing other fields + // Use explicit checks to allow 0 (zero cost models for budget bypass) + if ( + formValues.input_cost_per_token !== undefined && + formValues.input_cost_per_token !== null && + formValues.input_cost_per_token !== "" + ) { + formValues.input_cost_per_token = Number(formValues.input_cost_per_token) / 1000000; + } + if ( + formValues.output_cost_per_token !== undefined && + formValues.output_cost_per_token !== null && + formValues.output_cost_per_token !== "" + ) { + formValues.output_cost_per_token = Number(formValues.output_cost_per_token) / 1000000; + } + + // Cache Read Cost: if blank, default to Input Cost (already token-unit converted above) + if ( + formValues.cache_read_input_token_cost !== undefined && + formValues.cache_read_input_token_cost !== null && + formValues.cache_read_input_token_cost !== "" + ) { + formValues.cache_read_input_token_cost = Number(formValues.cache_read_input_token_cost) / 1000000; + } else if ( + formValues.input_cost_per_token !== undefined && + formValues.input_cost_per_token !== null && + formValues.input_cost_per_token !== "" + ) { + formValues.cache_read_input_token_cost = Number(formValues.input_cost_per_token); + } else { + delete formValues.cache_read_input_token_cost; + } + + // Cache Write Cost: explicit value if provided, else leave unset so the + // backend keeps the model-level default (per-second pricing, model_prices + // entries, etc.). Sending 0 here would overwrite that default. + // The backend falls back to input_cost_per_token when this key is absent. + if ( + formValues.cache_creation_input_token_cost !== undefined && + formValues.cache_creation_input_token_cost !== null && + formValues.cache_creation_input_token_cost !== "" + ) { + formValues.cache_creation_input_token_cost = Number(formValues.cache_creation_input_token_cost) / 1000000; + } else { + delete formValues.cache_creation_input_token_cost; + } + if ( + formValues.cache_creation_input_token_cost_above_1hr !== undefined && + formValues.cache_creation_input_token_cost_above_1hr !== null && + formValues.cache_creation_input_token_cost_above_1hr !== "" + ) { + formValues.cache_creation_input_token_cost_above_1hr = + Number(formValues.cache_creation_input_token_cost_above_1hr) / 1000000; + } else { + delete formValues.cache_creation_input_token_cost_above_1hr; + } + // Keep input_cost_per_second as is, no conversion needed + // Create a deployment for each mapping const deployments = []; for (const mapping of modelMappings) { @@ -34,55 +93,6 @@ export const prepareModelAddRequest = async (formValues: Record, ac const modelName = mapping.public_name; litellmParamsObj["model"] = mapping.litellm_model; - // Handle pricing conversion before processing other fields - // Use explicit checks to allow 0 (zero cost models for budget bypass) - if ( - formValues.input_cost_per_token !== undefined && - formValues.input_cost_per_token !== null && - formValues.input_cost_per_token !== "" - ) { - formValues.input_cost_per_token = Number(formValues.input_cost_per_token) / 1000000; - } - if ( - formValues.output_cost_per_token !== undefined && - formValues.output_cost_per_token !== null && - formValues.output_cost_per_token !== "" - ) { - formValues.output_cost_per_token = Number(formValues.output_cost_per_token) / 1000000; - } - - // Cache Read Cost: if blank, default to Input Cost (already token-unit converted above) - if ( - formValues.cache_read_input_token_cost !== undefined && - formValues.cache_read_input_token_cost !== null && - formValues.cache_read_input_token_cost !== "" - ) { - formValues.cache_read_input_token_cost = Number(formValues.cache_read_input_token_cost) / 1000000; - } else if ( - formValues.input_cost_per_token !== undefined && - formValues.input_cost_per_token !== null && - formValues.input_cost_per_token !== "" - ) { - formValues.cache_read_input_token_cost = Number(formValues.input_cost_per_token); - } else { - delete formValues.cache_read_input_token_cost; - } - - // Cache Write Cost: explicit value if provided, else leave unset so the - // backend keeps the model-level default (per-second pricing, model_prices - // entries, etc.). Sending 0 here would overwrite that default. - // The backend falls back to input_cost_per_token when this key is absent. - if ( - formValues.cache_creation_input_token_cost !== undefined && - formValues.cache_creation_input_token_cost !== null && - formValues.cache_creation_input_token_cost !== "" - ) { - formValues.cache_creation_input_token_cost = Number(formValues.cache_creation_input_token_cost) / 1000000; - } else { - delete formValues.cache_creation_input_token_cost; - } - // Keep input_cost_per_second as is, no conversion needed - // Iterate through the key-value pairs in formValues litellmParamsObj["model"] = mapping.litellm_model; for (const [key, value] of Object.entries(formValues)) { @@ -155,7 +165,8 @@ export const prepareModelAddRequest = async (formValues: Record, ac key === "output_cost_per_token" || key === "input_cost_per_second" || key === "cache_read_input_token_cost" || - key === "cache_creation_input_token_cost" + key === "cache_creation_input_token_cost" || + key === "cache_creation_input_token_cost_above_1hr" ) { if (value !== undefined && value !== null && value !== "") { litellmParamsObj[key] = Number(value); diff --git a/ui/litellm-dashboard/src/components/model_info_view.test.tsx b/ui/litellm-dashboard/src/components/model_info_view.test.tsx index 79717988292..db3e9354350 100644 --- a/ui/litellm-dashboard/src/components/model_info_view.test.tsx +++ b/ui/litellm-dashboard/src/components/model_info_view.test.tsx @@ -1695,13 +1695,54 @@ describe("ModelInfoView", () => { const [cacheRead, cacheWrite] = screen.getAllByPlaceholderText("Defaults to Input Cost if blank"); await user.type(cacheRead, "5"); await user.type(cacheWrite, "9"); + await user.type(screen.getByPlaceholderText("Defaults to the cost map 1 hour rate, else Cache Write Cost"), "15"); const payload = await save(user); - expect(payload.litellm_params).toMatchObject({ + const expectedPricing = { output_cost_per_token: 0.000012, cache_read_input_token_cost: 0.000005, cache_creation_input_token_cost: 0.000009, + cache_creation_input_token_cost_above_1hr: 0.000015, + }; + expect(payload.litellm_params).toMatchObject(expectedPricing); + }); + + describe("1 hour cache write tier", () => { + const withCacheWriteTiers = { + ...defaultModelData, + litellm_params: { + ...defaultModelData.litellm_params, + cache_creation_input_token_cost: 0.00000375, + cache_creation_input_token_cost_above_1hr: 0.000006, + }, + }; + + beforeEach(() => { + mockUseModelsInfo.mockReturnValue({ data: { data: [withCacheWriteTiers] }, isLoading: false, error: null }); + mockModelInfoV1Call.mockResolvedValue({ data: [withCacheWriteTiers] }); + }); + + it("shows both configured cache write rates per million tokens", async () => { + render(, { wrapper }); + + expect(await screen.findByText("Cache Write Cost (1 hour, per 1M tokens)")).toBeInTheDocument(); + expect(screen.getByText("3.7500")).toBeInTheDocument(); + expect(screen.getByText("6.0000")).toBeInTheDocument(); + }); + + it("clearing the 1 hour rate sends an explicit null and keeps the 5 minute rate", async () => { + const user = userEvent.setup(); + await enterEditMode(user); + + const oneHour = screen.getByPlaceholderText("Defaults to the cost map 1 hour rate, else Cache Write Cost"); + expect(oneHour).toHaveValue(6); + await user.clear(oneHour); + + const payload = await save(user); + + expect(payload.litellm_params.cache_creation_input_token_cost_above_1hr).toBeNull(); + expect(payload.litellm_params.cache_creation_input_token_cost).toBe(0.00000375); }); }); diff --git a/ui/litellm-dashboard/src/components/model_info_view.tsx b/ui/litellm-dashboard/src/components/model_info_view.tsx index 21fad41b2e5..a8df9999f52 100644 --- a/ui/litellm-dashboard/src/components/model_info_view.tsx +++ b/ui/litellm-dashboard/src/components/model_info_view.tsx @@ -51,6 +51,9 @@ import { Tag } from "./tag_management/types"; import { getDisplayModelName } from "./view_model/model_name_display"; import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog"; +const perTokenOrNull = (perMillion: string | number | null | undefined): number | null => + perMillion === undefined || perMillion === null || perMillion === "" ? null : Number(perMillion) / 1_000_000; + interface ModelInfoViewProps { modelId: string; onClose: () => void; @@ -327,19 +330,13 @@ export default function ModelInfoView({ } } - // Cache Write Cost: explicit value if provided, else explicit null so the - // backend removes the override and falls back to the model-level default. - // Sending 0 here would persist a zero rate even when the user intended to unset it. + // Cache write tiers: explicit value if provided, else explicit null so the backend removes the + // override. Sending 0 would persist a zero rate even when the user intended to unset it. if (isFieldTouched("cache_write_cost")) { - if ( - values.cache_write_cost !== undefined && - values.cache_write_cost !== null && - values.cache_write_cost !== "" - ) { - updatedLitellmParams.cache_creation_input_token_cost = Number(values.cache_write_cost) / 1_000_000; - } else { - updatedLitellmParams.cache_creation_input_token_cost = null; - } + updatedLitellmParams.cache_creation_input_token_cost = perTokenOrNull(values.cache_write_cost); + } + if (isFieldTouched("cache_write_cost_1h")) { + updatedLitellmParams.cache_creation_input_token_cost_above_1hr = perTokenOrNull(values.cache_write_cost_1h); } const storedCredentialName: string | null = localModelData?.litellm_params?.litellm_credential_name ?? null; diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index b77399c963f..bdf936d8bf3 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -49929,6 +49929,8 @@ export interface components { blocked?: boolean | null; /** Cache Creation Input Token Cost */ cache_creation_input_token_cost?: number | null; + /** Cache Creation Input Token Cost Above 1Hr */ + cache_creation_input_token_cost_above_1hr?: number | null; /** Cache Read Input Token Cost */ cache_read_input_token_cost?: number | null; /** Cost Per Ptu Per Hour */ From 60af2c601da29345ebd62df32819bff725934346 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 20:44:34 +0000 Subject: [PATCH 23/58] fix(proxy): emit spend queue size after updates and flushes (#46956) Co-authored-by: Mateo --- .../spend_update_queue.py | 3 + .../test_spend_update_queue.py | 69 +++++++++++++++++++ 2 files changed, 72 insertions(+) diff --git a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py index 2b8535cb113..e88b9fc2bfb 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py @@ -29,6 +29,7 @@ class SpendUpdateQueue(BaseUpdateQueue): ) -> DBSpendUpdateTransactions: """Flush all updates from the queue and return all updates aggregated by entity type.""" updates: Final = await self.flush_all_updates_from_in_memory_queue() + await self._emit_new_item_added_to_queue_event(queue_size=self.update_queue.qsize()) if len(updates) > 0: verbose_proxy_logger.info( "Spend tracking - flushed %d spend update items from in-memory queue", @@ -48,6 +49,7 @@ class SpendUpdateQueue(BaseUpdateQueue): "Spend update queue is full. Aggregating all entries in queue to concatenate entries." ) await self.aggregate_queue_updates() + await self._emit_new_item_added_to_queue_event(queue_size=self.update_queue.qsize()) async def aggregate_queue_updates(self): """Concatenate all updates in the queue to reduce the size of in-memory queue""" @@ -55,6 +57,7 @@ class SpendUpdateQueue(BaseUpdateQueue): aggregated_updates: Final = self._get_aggregated_spend_update_queue_item(updates) for update in aggregated_updates: await self.update_queue.put(update) + await self._emit_new_item_added_to_queue_event(queue_size=self.update_queue.qsize()) def _get_aggregated_spend_update_queue_item( self, updates: list[SpendUpdateQueueItem] diff --git a/tests/unit/proxy/db/db_transaction_queue/test_spend_update_queue.py b/tests/unit/proxy/db/db_transaction_queue/test_spend_update_queue.py index 43f1a820885..9a4b02ba3ba 100644 --- a/tests/unit/proxy/db/db_transaction_queue/test_spend_update_queue.py +++ b/tests/unit/proxy/db/db_transaction_queue/test_spend_update_queue.py @@ -1,12 +1,81 @@ import asyncio import json +from typing import Final import pytest from fastapi.testclient import TestClient +from prometheus_client import REGISTRY +import litellm from litellm.constants import MAX_SIZE_IN_MEMORY_QUEUE from litellm.proxy._types import Litellm_EntityType, SpendUpdateQueueItem from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue +from litellm.types.services import ServiceTypes + + +@pytest.mark.asyncio +@pytest.mark.parametrize("aggregation_threshold", (2, 3)) +async def test_spend_queue_gauge_tracks_enqueue_aggregation_and_drain( + monkeypatch: pytest.MonkeyPatch, aggregation_threshold: int +) -> None: + monkeypatch.setattr(litellm, "service_callback", ["prometheus_system"]) + queue: Final = SpendUpdateQueue() + queue.MAX_SIZE_IN_MEMORY_QUEUE = aggregation_threshold + response_cost: Final = 0.5 + update: Final = SpendUpdateQueueItem( + entity_type=Litellm_EntityType.USER, entity_id="queue-user", response_cost=response_cost + ) + + await queue.add_update(update) + await asyncio.sleep(0) + assert REGISTRY.get_sample_value( + "litellm_in_memory_spend_update_queue_size", + {"in_memory_spend_update_queue": str(ServiceTypes.IN_MEMORY_SPEND_UPDATE_QUEUE)}, + ) == queue.update_queue.qsize() == 1 + + await queue.add_update(update) + await asyncio.sleep(0) + assert REGISTRY.get_sample_value( + "litellm_in_memory_spend_update_queue_size", + {"in_memory_spend_update_queue": str(ServiceTypes.IN_MEMORY_SPEND_UPDATE_QUEUE)}, + ) == queue.update_queue.qsize() == (1 if aggregation_threshold == 2 else 2) + + await queue.aggregate_queue_updates() + await asyncio.sleep(0) + assert REGISTRY.get_sample_value( + "litellm_in_memory_spend_update_queue_size", + {"in_memory_spend_update_queue": str(ServiceTypes.IN_MEMORY_SPEND_UPDATE_QUEUE)}, + ) == queue.update_queue.qsize() == 1 + + transactions: Final = await queue.flush_and_get_aggregated_db_spend_update_transactions() + await asyncio.sleep(0) + assert transactions["user_list_transactions"] == {"queue-user": 2 * response_cost} + assert REGISTRY.get_sample_value( + "litellm_in_memory_spend_update_queue_size", + {"in_memory_spend_update_queue": str(ServiceTypes.IN_MEMORY_SPEND_UPDATE_QUEUE)}, + ) == queue.update_queue.qsize() == 0 + + +@pytest.mark.asyncio +async def test_spend_queue_gauge_tracks_partial_flush(monkeypatch: pytest.MonkeyPatch) -> None: + import litellm.proxy.db.db_transaction_queue.base_update_queue as base_queue_module + + monkeypatch.setattr(litellm, "service_callback", ["prometheus_system"]) + monkeypatch.setattr(base_queue_module, "MAX_IN_MEMORY_QUEUE_FLUSH_COUNT", 1) + queue: Final = SpendUpdateQueue() + update: Final = SpendUpdateQueueItem( + entity_type=Litellm_EntityType.USER, entity_id="queue-user", response_cost=0.5 + ) + await queue.add_update(update) + await queue.add_update(update) + + transactions: Final = await queue.flush_and_get_aggregated_db_spend_update_transactions() + await asyncio.sleep(0) + assert transactions["user_list_transactions"] == {"queue-user": update["response_cost"]} + assert REGISTRY.get_sample_value( + "litellm_in_memory_spend_update_queue_size", + {"in_memory_spend_update_queue": str(ServiceTypes.IN_MEMORY_SPEND_UPDATE_QUEUE)}, + ) == queue.update_queue.qsize() == 1 From 388251a5b863cf1da7b1829875a735f0a398fa9a Mon Sep 17 00:00:00 2001 From: AxelRay Date: Sat, 17 Oct 2026 04:00:37 +0700 Subject: [PATCH 24/58] fix(bedrock): drop dangerous-tool-use beta without safeguards (#44471) * fix(bedrock): drop dangerous beta without safeguards * fix(bedrock): drop the dangerous-tool-use beta only next to structured outputs, keep it on Mantle Bedrock rejects the pair with 'invalid beta flag' on Claude models before Opus 4.7, while requests without structured outputs and Bedrock Mantle accept the beta, so the drop now only applies to Invoke requests with output_config.format or a strict tool and no safeguards. * fix(anthropic): keep the AnthropicThinkingParam import that main now uses * fix(bedrock): drop the client's dangerous-tool-use beta whenever safeguards is absent Bedrock rejects the beta with 'invalid beta flag' on Sonnet 4.5, Opus 4.5, and Opus 4.6 even without structured outputs, so dropping it only next to structured outputs left those beta-only requests failing. The beta does nothing without safeguards, so Invoke now drops it in every such request. Mantle keeps it. The thinking display helper goes back to main's version. * fix(bedrock): drop the client's dangerous-tool-use beta on Invoke chat without safeguards * test(bedrock): wire cells for the dangerous-tool-use beta on Invoke messages, chat and responses * test(bedrock): type the new dangerous-tool-use beta tests * test(bedrock_mantle): type the dangerous-tool-use beta Mantle test --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../anthropic_claude3_transformation.py | 8 +- .../anthropic_claude3_transformation.py | 14 +- .../bedrock/messages/mantle_transformation.py | 3 + ...ock_invoke_dangerous_tool_use_beta_wire.py | 138 ++++++++++++++++++ ...ations_anthropic_claude3_transformation.py | 29 ++++ .../test_anthropic_claude3_transformation.py | 54 +++++++ ..._bedrock_mantle_messages_transformation.py | 9 ++ 7 files changed, 253 insertions(+), 2 deletions(-) create mode 100644 tests/integration/providers/test_bedrock_invoke_dangerous_tool_use_beta_wire.py diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index 04747405375..fc6088ddc66 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -32,6 +32,7 @@ from litellm.llms.bedrock.common_utils import ( tools_without_eager_input_streaming, ) from litellm.types.llms.anthropic import ( + ANTHROPIC_BETA_HEADER_VALUES, ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER, ANTHROPIC_TOOL_SEARCH_BETA_HEADER, AnthropicThinkingParam, @@ -270,7 +271,12 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): is_mid_conversation_output_config_used: Final = self.is_mid_conversation_output_config_used(messages) is_thinking_display_updates_used: Final = self.is_thinking_display_updates_used(thinking) - user_beta_set: Final = set(get_anthropic_beta_from_headers(headers)) + client_beta_set: Final = set(get_anthropic_beta_from_headers(headers)) + user_beta_set: Final = ( + client_beta_set + if "safeguards" in optional_params + else client_beta_set - {ANTHROPIC_BETA_HEADER_VALUES.DANGEROUS_TOOL_USE_2026_09_03.value} + ) beta_set: Final = set(user_beta_set) auto_betas: Final = self.get_anthropic_beta_list( model=model, diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 88ae0d4d514..95e81fa3ee7 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -90,6 +90,9 @@ class AmazonAnthropicClaudeMessagesConfig( def beta_headers_provider(self) -> str: return self.custom_llm_provider or "bedrock" + def drops_dangerous_tool_use_beta_without_safeguards(self) -> bool: + return True + BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS = frozenset(BedrockInvokeAnthropicMessagesRequest.__annotations__.keys()) def get_error_class( @@ -552,8 +555,17 @@ class AmazonAnthropicClaudeMessagesConfig( if anthropic_model_info.is_eager_input_streaming_used(tools): beta_set.add(ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER) + dangerous_tool_use_beta: Final = ANTHROPIC_BETA_HEADER_VALUES.DANGEROUS_TOOL_USE_2026_09_03.value if anthropic_messages_optional_request_params.get("safeguards") is not None: - beta_set.add(ANTHROPIC_BETA_HEADER_VALUES.DANGEROUS_TOOL_USE_2026_09_03.value) + beta_set.add(dangerous_tool_use_beta) + elif dangerous_tool_use_beta in beta_set and self.drops_dangerous_tool_use_beta_without_safeguards(): + beta_set.discard(dangerous_tool_use_beta) + verbose_logger.debug( + "Bedrock Invoke: dropping %s from client headers because the request has no safeguards. Bedrock " + "answers it with 'invalid beta flag' next to structured outputs before Opus 4.7, and on Sonnet 4.5, " + "Opus 4.5, and Opus 4.6 even alone.", + dangerous_tool_use_beta, + ) self._filter_context_management_for_bedrock_invoke( anthropic_messages_request=anthropic_messages_request, diff --git a/litellm/llms/bedrock/messages/mantle_transformation.py b/litellm/llms/bedrock/messages/mantle_transformation.py index ca93cc3e617..1ad84761bc9 100644 --- a/litellm/llms/bedrock/messages/mantle_transformation.py +++ b/litellm/llms/bedrock/messages/mantle_transformation.py @@ -67,6 +67,9 @@ class AmazonMantleMessagesConfig(AmazonAnthropicClaudeMessagesConfig): def should_filter_anthropic_beta_headers(self) -> bool: return False + def drops_dangerous_tool_use_beta_without_safeguards(self) -> bool: + return False + def get_complete_url( self, api_base: str | None, diff --git a/tests/integration/providers/test_bedrock_invoke_dangerous_tool_use_beta_wire.py b/tests/integration/providers/test_bedrock_invoke_dangerous_tool_use_beta_wire.py new file mode 100644 index 00000000000..ca22ab4f23d --- /dev/null +++ b/tests/integration/providers/test_bedrock_invoke_dangerous_tool_use_beta_wire.py @@ -0,0 +1,138 @@ +import json +from collections.abc import Callable, Mapping +from typing import Final +from urllib.parse import unquote + +import pytest +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +CLAUDE_ID: Final = "us.anthropic.claude-sonnet-4-5-20250929-v1:0" +MANTLE_CLAUDE_ID: Final = "anthropic.claude-haiku-4-5" +TOKEN: Final = "synthetic-bedrock-bearer" +DANGEROUS_TOOL_USE_BETA: Final = "dangerous-tool-use-2026-09-03" +KEPT_BETA: Final = "context-1m-2025-08-07" +SAFEGUARDS: Final[JsonValue] = [{"type": "dangerous_tool_use"}] +NO_CACHE: Final[JsonValue] = {"no-cache": True} +INVOKE_TARGET: Final = f"/model/{CLAUDE_ID}/invoke" +MANTLE_TARGET: Final = "/anthropic/v1/messages" +REPLY: Final = json.dumps( + { + "id": "msg_scripted", + "type": "message", + "role": "assistant", + "model": CLAUDE_ID, + "content": [{"type": "text", "text": "scripted reply"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 3, "output_tokens": 2}, + } +).encode() +_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_BETAS: Final = TypeAdapter(tuple[str, ...]) + + +def _peer(request: Request) -> Reply: + if unquote(request.target) not in (INVOKE_TARGET, MANTLE_TARGET): + return Reply(status=404, body=json.dumps({"error": f"unscripted target {request.target}"}).encode()) + assert request.headers["authorization"] == f"Bearer {TOKEN}", request.headers + return Reply(body=REPLY) + + +def _invoke_body_betas(request: Request) -> frozenset[str]: + return frozenset(_BETAS.validate_python(_OBJECT.validate_json(request.body).get("anthropic_beta") or ())) + + +def _mantle_header_betas(request: Request) -> frozenset[str]: + return frozenset(beta.strip() for beta in request.headers.get("anthropic-beta", "").split(",") if beta.strip()) + + +def _delivered(gateway: Gateway, wire: Wire, path: str, body: Mapping[str, JsonValue], client_betas: str) -> Request: + response: Final = gateway.request("POST", path, body, headers={"anthropic-beta": client_betas}) + assert response.status_code == 200, response.text + (sent,) = wire.drain() + return sent + + +def _messages_body(model: str, **extra: JsonValue) -> dict[str, JsonValue]: + return { + "model": model, + "max_tokens": 16, + "messages": [{"role": "user", "content": "hi"}], + "cache": NO_CACHE, + **extra, + } + + +def _chat_body(model: str, **extra: JsonValue) -> dict[str, JsonValue]: + return {"model": model, "messages": [{"role": "user", "content": "hi"}], "cache": NO_CACHE, **extra} + + +def _responses_body(model: str, **extra: JsonValue) -> dict[str, JsonValue]: + return {"model": model, "input": "hi", "cache": NO_CACHE, **extra} + + +def test_invoke_messages_without_safeguards_sends_the_betas_of_a_request_that_never_had_the_beta( + gateway: Gateway, +) -> None: + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"bedrock/{CLAUDE_ID}", api_key=TOKEN, aws_region_name="us-east-1", api_base=wire.url + ) + with_beta: Final = _invoke_body_betas( + _delivered(gateway, wire, "/v1/messages", _messages_body(model), f"{KEPT_BETA},{DANGEROUS_TOOL_USE_BETA}") + ) + without_beta: Final = _invoke_body_betas( + _delivered(gateway, wire, "/v1/messages", _messages_body(model), KEPT_BETA) + ) + assert KEPT_BETA in without_beta, without_beta + assert with_beta == without_beta, (with_beta, without_beta) + + +def test_invoke_messages_with_safeguards_keeps_the_dangerous_beta(gateway: Gateway) -> None: + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"bedrock/{CLAUDE_ID}", api_key=TOKEN, aws_region_name="us-east-1", api_base=wire.url + ) + sent: Final = _invoke_body_betas( + _delivered( + gateway, + wire, + "/v1/messages", + _messages_body(model, safeguards=SAFEGUARDS), + f"{KEPT_BETA},{DANGEROUS_TOOL_USE_BETA}", + ) + ) + assert {KEPT_BETA, DANGEROUS_TOOL_USE_BETA} <= sent, sent + + +@pytest.mark.parametrize( + ("path", "build"), + [("/v1/chat/completions", _chat_body), ("/v1/responses", _responses_body)], + ids=["chat_completions", "responses"], +) +def test_invoke_chat_shaped_routes_without_safeguards_drop_only_the_dangerous_beta( + gateway: Gateway, path: str, build: Callable[[str], dict[str, JsonValue]] +) -> None: + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"bedrock/invoke/{CLAUDE_ID}", api_key=TOKEN, aws_region_name="us-east-1", api_base=wire.url + ) + with_beta: Final = _invoke_body_betas( + _delivered(gateway, wire, path, build(model), f"{KEPT_BETA},{DANGEROUS_TOOL_USE_BETA}") + ) + without_beta: Final = _invoke_body_betas(_delivered(gateway, wire, path, build(model), KEPT_BETA)) + assert KEPT_BETA in without_beta, without_beta + assert with_beta == without_beta, (with_beta, without_beta) + + +def test_mantle_messages_without_safeguards_keeps_the_dangerous_beta(gateway: Gateway) -> None: + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"bedrock_mantle/{MANTLE_CLAUDE_ID}", api_key=TOKEN, aws_region_name="us-east-1", api_base=wire.url + ) + sent: Final = _mantle_header_betas( + _delivered(gateway, wire, "/v1/messages", _messages_body(model), f"{KEPT_BETA},{DANGEROUS_TOOL_USE_BETA}") + ) + assert DANGEROUS_TOOL_USE_BETA in sent, sent diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index 72857246d85..16f44f4f35f 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -1087,6 +1087,35 @@ def test_bedrock_chat_invoke_eager_input_streaming_beta_not_duplicated_with_clie assert result["anthropic_beta"] == [FINE_GRAINED_TOOL_STREAMING_BETA] +DANGEROUS_TOOL_USE_BETA: Final = "dangerous-tool-use-2026-09-03" + + +def _chat_invoke_betas(extra_optional_params: dict[str, object]) -> object: + config: Final = AmazonAnthropicClaudeConfig() + model: Final = "us.anthropic.claude-sonnet-4-5-20250929-v1:0" + optional_params: Final = config.map_openai_params( + non_default_params={"max_tokens": 64}, optional_params={}, model=model, drop_params=False + ) + return config.transform_request( + model=model, + messages=[{"role": "user", "content": "hi"}], + optional_params={**optional_params, **extra_optional_params}, + litellm_params={}, + headers={"anthropic-beta": f"claude-code-20250219,{DANGEROUS_TOOL_USE_BETA}"}, + ).get("anthropic_beta") + + +def test_bedrock_chat_invoke_drops_client_dangerous_tool_use_beta_without_safeguards() -> None: + assert _chat_invoke_betas({}) == ["claude-code-20250219"] + + +def test_bedrock_chat_invoke_keeps_client_dangerous_tool_use_beta_with_safeguards() -> None: + assert _chat_invoke_betas({"safeguards": [{"type": "dangerous_tool_use"}]}) == [ + "claude-code-20250219", + DANGEROUS_TOOL_USE_BETA, + ] + + def _mid_conversation_system_conversation() -> list[dict]: return [ { diff --git a/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index dffe2c36e20..c91eb75dbd9 100644 --- a/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -1677,6 +1677,60 @@ def test_bedrock_messages_does_not_add_dangerous_tool_use_beta_without_safeguard assert "dangerous-tool-use-2026-09-03" not in result.get("anthropic_beta", []) +_NAME_SCHEMA: Final = {"type": "object", "properties": {"name": {"type": "string"}}, "required": ["name"]} + + +@pytest.mark.parametrize( + "request_params", + [ + {}, + {"output_config": {"format": {"type": "json_schema", "schema": _NAME_SCHEMA}}}, + {"tools": [{"name": "get_name", "strict": True, "input_schema": _NAME_SCHEMA}]}, + ], + ids=["plain", "output_config_format", "strict_tool"], +) +def test_bedrock_messages_drops_client_dangerous_tool_use_beta_without_safeguards( + local_model_cost_map: None, local_beta_headers_config: None, request_params: dict[str, object] +) -> None: + from litellm.types.router import GenericLiteLLMParams + + result: Final = AmazonAnthropicClaudeMessagesConfig().transform_anthropic_messages_request( + model="us.anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": [{"type": "text", "text": "Name a color."}]}], + anthropic_messages_optional_request_params={"max_tokens": 64, **request_params}, + litellm_params=GenericLiteLLMParams(), + headers={"anthropic-beta": "dangerous-tool-use-2026-09-03"}, + ) + + assert all(result.get(key) == value for key, value in request_params.items()) + assert "dangerous-tool-use-2026-09-03" not in result.get("anthropic_beta", []) + + +def test_bedrock_messages_keeps_client_dangerous_tool_use_beta_with_safeguards_and_structured_outputs( + local_model_cost_map: None, local_beta_headers_config: None +) -> None: + from litellm.types.router import GenericLiteLLMParams + + safeguards: Final = [{"type": "dangerous_tool_use", "classifier_context": {"v": 1, "permission_mode": "auto"}}] + strict_tool: Final = {"name": "get_name", "strict": True, "input_schema": _NAME_SCHEMA} + + result: Final = AmazonAnthropicClaudeMessagesConfig().transform_anthropic_messages_request( + model="us.anthropic.claude-sonnet-5", + messages=[{"role": "user", "content": [{"type": "text", "text": "Name a color."}]}], + anthropic_messages_optional_request_params={ + "max_tokens": 64, + "safeguards": safeguards, + "tools": [strict_tool], + }, + litellm_params=GenericLiteLLMParams(), + headers={"anthropic-beta": "dangerous-tool-use-2026-09-03"}, + ) + + assert result["safeguards"] == safeguards + assert result["tools"] == [strict_tool] + assert result["anthropic_beta"].count("dangerous-tool-use-2026-09-03") == 1 + + def test_bedrock_messages_stream_decoder_keeps_safeguard_results(): """Bedrock streams the classifier verdicts on message_start and on the final message_delta, exactly as api.anthropic.com does.""" decoder = AmazonAnthropicClaudeMessagesStreamDecoder(model="us.anthropic.claude-sonnet-5") diff --git a/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py b/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py index 923572c4f46..3236efbbfba 100644 --- a/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py +++ b/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py @@ -10,6 +10,7 @@ Bearer/SigV4 auth chain, and the wire request through the public entrypoint. """ import json +from typing import Final from unittest.mock import MagicMock import httpx @@ -453,6 +454,14 @@ class TestBetaHeadersOnTheWire: assert _sent_betas(route) == ["dangerous-tool-use-2026-09-03"] assert _sent_body(route)["safeguards"] == safeguards + @pytest.mark.asyncio + @respx.mock + async def test_client_dangerous_tool_use_beta_reaches_mantle_without_safeguards(self) -> None: + route: Final = await self._send(extra_headers={"anthropic-beta": "dangerous-tool-use-2026-09-03"}) + + assert _sent_betas(route) == ["dangerous-tool-use-2026-09-03"] + assert "safeguards" not in _sent_body(route) + @pytest.mark.asyncio @respx.mock async def test_betas_and_version_never_travel_in_the_body(self): From 78372b440d7c353aeb648926fbf3c69106cb8a4b Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Fri, 16 Oct 2026 14:03:10 -0700 Subject: [PATCH 25/58] feat(cost-map): add together_ai reflection/Beam-1.0-501B-A23B (#46964) Price-Sync: litellm-providers Co-authored-by: berriai-litellm-provider-info-sync[bot] <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 13 +++++++++++++ model_prices_and_context_window.json | 13 +++++++++++++ 2 files changed, 26 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 49d2ce90af4..15bf988c96b 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -75218,5 +75218,18 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true + }, + "together_ai/reflection/Beam-1.0-501B-A23B": { + "cache_read_input_token_cost": 4e-08, + "input_cost_per_token": 4e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://api.together.ai/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true } } diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 49d2ce90af4..15bf988c96b 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -75218,5 +75218,18 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true + }, + "together_ai/reflection/Beam-1.0-501B-A23B": { + "cache_read_input_token_cost": 4e-08, + "input_cost_per_token": 4e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://api.together.ai/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true } } From 5bcced0187d8091c4cb21553c6c18b7cabfa3f41 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 16 Oct 2026 14:04:36 -0700 Subject: [PATCH 26/58] feat(ui): redesign the playground chat around a single composer and inspector (#46934) * feat(ui): add clean-slate playground chat prototype * feat(ui): replace the playground chat with the redesigned composer and inspector * fix(ui): close the playground chat gaps found in self-review Resend only complete text exchanges, so empty replies, MCP direct calls and transcriptions no longer leak into chat history. A stopped request can no longer append output, attach its gateway trace to a newer reply or clear the newer request's loading state. MCP selections are only pruned once both the server and toolset lists load. The MCP tool picker lists toolsets again and calls the server that backs the chosen tool, and a saved model the current key cannot use is cleared when the model list loads, matching the old chat Also vendors the scroll-fade-x utility the attachment primitive uses * fix(ui): resolve playground toolset tools per server and drop late session ids A toolset now lists only its own tools, each resolved on the server the toolset names, so the args form and the call always target the same tool. A saved toolset no longer asks for tools with its own id while the toolset list loads, and a failed tool list says so instead of claiming the server has no tools A stopped request can no longer overwrite the Responses session id, starting a new chat drops a request that is still running, embeddings replies are no longer resent as chat history, and a saved model missing from an already cached model list is cleared too. A failing toolset list falls back right away instead of retrying while the MCP pickers show loading * fix(ui): guard playground sends and label inspector stats Pressing Enter twice no longer sends the same message twice, an attachment the new endpoint cannot take is dropped when the endpoint changes, a stopped Responses request can no longer show its code output on a newer reply, and a reply that broke mid-stream is no longer resent as history The inspector's speed, cost and routing stats get the same icons as the reply footer, and the footer's metrics now line up with its copy and Inspect buttons * fix(ui): write playground key settings only through the encoded storage path Key source, virtual key and custom headers now go through their own write loop, separate from the plain settings, so a key value can never take the plain sessionStorage branch * fix(ui): restore playground BYOK connect and reconnect Picking a BYOK server for a direct tool call now shows a key button next to the server picker, so the key can be connected without leaving the MCP endpoint. A connected BYOK server gets a Replace key action in the MCP dialog, also shown when its tools fail to load, and saving a key reloads that server's tool list * revert: restore the single playground settings storage loop Reverts ca4e742115. The CodeQL finding covers the same encoded sessionStorage handling the old chat used, so it is left as is * test(ui): type the BYOK credential mock without deep instantiation * fix(ui): let users back out of replacing a BYOK key The replace-key form in the MCP dialog now has Cancel, and it resets when another server is picked or the dialog closes, so a working key is never forced to be overwritten * refactor(ui): name the MCP dialog row focus handler * fix(ui): accept exact custom models, tags and decimal temperatures in the playground A typed model or tag that is a prefix of a listed one could not be used, because Enter picked the partial match and the custom action only showed when nothing matched. The custom value is now a list option whenever no exact match exists. Temperature gets an editable input beside the slider so values finer than the 0.1 step can be sent again. * fix(ui): stop claiming no guardrails ran when the gateway sent no headers Direct MCP calls, A2A agents and any request that fails before headers arrive built their trace from an empty header set, so the inspector said no guardrails or policies ran. The gateway fields are now left unset in that case and the inspector says the endpoint does not report them. * fix(ui): leave policies unknown when the gateway headers are hidden The proxy exposes x-litellm-applied-guardrails to other origins but not x-litellm-applied-policies, so a dashboard on another origin read every request as having run no policies. Policies are now unknown unless the call id header is readable, and the inspector says so. --- tests/e2e/ui/helpers/playground.ts | 55 +- .../guardrails/presidioUserStory.spec.ts | 10 +- .../tests/internal-user/internalUser.spec.ts | 9 +- .../ui/tests/playground/playground.spec.ts | 23 +- .../e2e/ui/tests/team-admin/teamAdmin.spec.ts | 15 +- ui/litellm-dashboard/eslint-suppressions.json | 15 + .../components/chat/ByokConnect.tsx | 152 ++++ .../components/chat/Composer.test.tsx | 61 ++ .../playground/components/chat/Composer.tsx | 196 ++++++ .../components/chat/ComposerControls.tsx | 120 ++++ .../components/chat/Conversation.tsx | 388 +++++++++++ .../playground/components/chat/Inspector.tsx | 292 ++++++++ .../playground/components/chat/McpDialog.tsx | 391 +++++++++++ .../playground/components/chat/ModelName.tsx | 74 ++ .../components/chat/Pickers.test.tsx | 83 +++ .../playground/components/chat/Pickers.tsx | 193 +++++ .../chat/PlaygroundChat.integration.test.tsx | 658 ++++++++++++++++++ .../components/chat/PlaygroundChat.tsx | 489 +++++++++++++ .../components/chat/ReplyContent.tsx | 92 +++ .../components/chat/SettingsPopover.tsx | 395 +++++++++++ .../components/chat/composerFiles.ts | 52 ++ .../components/chat/dispatchRequest.test.ts | 58 ++ .../components/chat/dispatchRequest.ts | 295 ++++++++ .../components/chat/endpointCapabilities.ts | 104 +++ .../chat/playgroundSettings.test.ts | 149 ++++ .../components/chat/playgroundSettings.ts | 192 +++++ .../playground/components/chat/useIsWide.ts | 16 + .../components/chat/usePlaygroundCatalog.ts | 147 ++++ .../components/chat/usePlaygroundSend.ts | 263 +++++++ .../components/chat/usePlaygroundSettings.ts | 38 + .../llm_calls/anthropic_messages.tsx | 3 + .../playground/llm_calls/audio_speech.tsx | 3 + .../llm_calls/audio_transcriptions.tsx | 3 + .../playground/llm_calls/embeddings_api.tsx | 3 + .../playground/llm_calls/image_edits.tsx | 3 + .../playground/llm_calls/image_generation.tsx | 3 + .../playground/llm_calls/interactions_api.tsx | 3 + .../app/(dashboard)/playground/page.test.tsx | 4 +- .../src/app/(dashboard)/playground/page.tsx | 87 ++- ui/litellm-dashboard/src/app/globals.css | 196 ++++++ .../src/components/chat_ui/types.ts | 24 + .../components/llm_calls/chat_completion.tsx | 23 +- .../components/llm_calls/reply_trace.test.ts | 64 ++ .../src/components/llm_calls/reply_trace.ts | 49 ++ .../components/llm_calls/responses_api.tsx | 18 +- .../src/components/ui/accordion.tsx | 59 ++ .../src/components/ui/attachment.tsx | 185 +++++ .../src/components/ui/item.tsx | 179 +++++ .../src/components/ui/kbd.tsx | 20 + .../src/components/ui/marker.tsx | 67 ++ .../src/components/ui/message-scroller.tsx | 117 ++++ .../src/components/ui/message.tsx | 74 ++ .../src/components/ui/spinner.tsx | 16 + .../src/components/ui/toggle-group.tsx | 87 +++ .../src/components/ui/toggle.tsx | 37 + .../src/lib/http/observeResponseHeaders.ts | 9 + 56 files changed, 6245 insertions(+), 116 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/ByokConnect.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/Composer.test.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/Composer.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/ComposerControls.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/Conversation.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/Inspector.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/McpDialog.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/ModelName.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/Pickers.test.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/Pickers.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/PlaygroundChat.integration.test.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/PlaygroundChat.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/ReplyContent.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/SettingsPopover.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/composerFiles.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/dispatchRequest.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/dispatchRequest.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/endpointCapabilities.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/playgroundSettings.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/playgroundSettings.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/useIsWide.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/usePlaygroundCatalog.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/usePlaygroundSend.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/usePlaygroundSettings.ts create mode 100644 ui/litellm-dashboard/src/components/llm_calls/reply_trace.test.ts create mode 100644 ui/litellm-dashboard/src/components/llm_calls/reply_trace.ts create mode 100644 ui/litellm-dashboard/src/components/ui/accordion.tsx create mode 100644 ui/litellm-dashboard/src/components/ui/attachment.tsx create mode 100644 ui/litellm-dashboard/src/components/ui/item.tsx create mode 100644 ui/litellm-dashboard/src/components/ui/kbd.tsx create mode 100644 ui/litellm-dashboard/src/components/ui/marker.tsx create mode 100644 ui/litellm-dashboard/src/components/ui/message-scroller.tsx create mode 100644 ui/litellm-dashboard/src/components/ui/message.tsx create mode 100644 ui/litellm-dashboard/src/components/ui/spinner.tsx create mode 100644 ui/litellm-dashboard/src/components/ui/toggle-group.tsx create mode 100644 ui/litellm-dashboard/src/components/ui/toggle.tsx create mode 100644 ui/litellm-dashboard/src/lib/http/observeResponseHeaders.ts diff --git a/tests/e2e/ui/helpers/playground.ts b/tests/e2e/ui/helpers/playground.ts index 866334e1c94..b46a2263072 100644 --- a/tests/e2e/ui/helpers/playground.ts +++ b/tests/e2e/ui/helpers/playground.ts @@ -2,42 +2,57 @@ import { expect, type Locator, type Page as PlaywrightPage } from "@playwright/t import { navigateToPage, dismissFeedbackPopup } from "./navigation"; import { Page } from "../fixtures/pages"; -/** Controls for the Test Key / Playground page, shared with the router-fallback specs. */ - -/** - * The configuration panel is rendered twice, docked and overlay, with one visible at a time. - * Every control is narrowed to the visible copy or it trips strict mode against its hidden twin. - */ export const onlyVisible = (locator: Locator): Locator => locator.filter({ visible: true }).first(); -/** The model combobox, addressed by the placeholder its search input shows before selection. */ -export const modelSelect = (page: PlaywrightPage): Locator => onlyVisible(page.getByPlaceholder("Select a Model")); +export const composer = (page: PlaywrightPage): Locator => + onlyVisible(page.getByRole("group").filter({ has: page.getByRole("textbox", { name: "Message" }) })); + +export const modelSelect = (page: PlaywrightPage): Locator => + onlyVisible(page.getByRole("combobox", { name: "Model", exact: true })); export const sendButton = (page: PlaywrightPage): Locator => - onlyVisible(page.getByRole("button", { name: "Send message" })); + composer(page).getByRole("button", { name: "Send message" }); -/** The Virtual Key Source dropdown, addressed by the accessible name on its trigger. */ -export const keySourceSelect = (page: PlaywrightPage): Locator => - onlyVisible(page.getByLabel("Virtual Key Source")); +export const keySourceSelect = (page: PlaywrightPage): Locator => onlyVisible(page.getByLabel("Send as")); export async function openPlayground(page: PlaywrightPage): Promise { await navigateToPage(page, Page.LlmPlayground); await dismissFeedbackPopup(page); - await expect(onlyVisible(page.getByText("Virtual Key Source"))).toBeVisible({ - timeout: 20_000, - }); + await expect(modelSelect(page)).toBeVisible({ timeout: 20_000 }); + const searchHint = page.getByRole("button", { name: "Dismiss search hint" }); + if (await searchHint.isVisible()) { + await searchHint.click(); + } +} + +export async function openRequestSettings(page: PlaywrightPage): Promise { + await composer(page).getByRole("button", { name: "Settings", exact: true }).click(); + await expect(keySourceSelect(page)).toBeVisible({ timeout: 10_000 }); +} + +export async function enterVirtualKey(page: PlaywrightPage, key: string): Promise { + await openRequestSettings(page); + await keySourceSelect(page).click(); + await onlyVisible(page.getByRole("option", { name: "Virtual key", exact: true })).click({ timeout: 15_000 }); + const keyInput = onlyVisible(page.getByLabel("Virtual key", { exact: true })); + await expect(keyInput).toBeVisible({ timeout: 10_000 }); + await keyInput.fill(key); + await page.keyboard.press("Escape"); +} + +export async function searchModels(page: PlaywrightPage, query: string): Promise { + await modelSelect(page).click(); + await onlyVisible(page.getByPlaceholder("Search models")).fill(query); } export async function selectModel(page: PlaywrightPage, model: string): Promise { - const select = modelSelect(page); - await select.click(); - // Virtualized: options outside the rendered window are absent from the DOM, so search first. - await select.fill(model); + await searchModels(page, model); await onlyVisible(page.getByRole("option", { name: model, exact: true })).click({ timeout: 15_000 }); + await expect(modelSelect(page)).toContainText(model); } export async function sendMessage(page: PlaywrightPage, message: string): Promise { - const input = onlyVisible(page.getByPlaceholder("Type your message", { exact: false })); + const input = composer(page).getByRole("textbox", { name: "Message" }); await expect(input).toBeVisible({ timeout: 15_000 }); await input.fill(message); await sendButton(page).click(); diff --git a/tests/e2e/ui/tests/guardrails/presidioUserStory.spec.ts b/tests/e2e/ui/tests/guardrails/presidioUserStory.spec.ts index b4dc7c4e8da..c57a9147ab0 100644 --- a/tests/e2e/ui/tests/guardrails/presidioUserStory.spec.ts +++ b/tests/e2e/ui/tests/guardrails/presidioUserStory.spec.ts @@ -3,7 +3,7 @@ import { ADMIN_STORAGE_PATH, MOCK_PRESIDIO_URL, PROPAGATION_TIMEOUT_MS } from ". import { navigateToPage, dismissFeedbackPopup } from "../../helpers/navigation"; import { Page } from "../../fixtures/pages"; import { CHAT_MODEL_A, masterKey, rootPath, uniqueSuffix, waitForSpendLog } from "../../helpers/traffic"; -import { openPlayground, selectModel, sendButton, onlyVisible } from "../../helpers/playground"; +import { composer, openPlayground, selectModel, sendButton, onlyVisible } from "../../helpers/playground"; const RAW_EMAIL = "jane.doe@example.com"; const RAW_PHONE = "555-867-5309"; @@ -123,14 +123,12 @@ test.describe("Presidio PII guardrail, end to end from the dashboard", () => { await openPlayground(page); await selectModel(page, CHAT_MODEL_A); - const guardrailSelect = onlyVisible(page.getByPlaceholder("Select guardrails")); - await expect(guardrailSelect).toBeVisible({ timeout: 20_000 }); - await guardrailSelect.click(); - await guardrailSelect.fill(guardrailName); + await composer(page).getByRole("combobox", { name: "Guardrails" }).click(); + await onlyVisible(page.getByPlaceholder("Search guardrails")).fill(guardrailName); await onlyVisible(page.getByRole("option", { name: guardrailName, exact: true })).click({ timeout: 20_000 }); await page.keyboard.press("Escape"); - const input = onlyVisible(page.getByPlaceholder("Type your message", { exact: false })); + const input = composer(page).getByRole("textbox", { name: "Message" }); await expect(input).toBeVisible({ timeout: 15_000 }); await input.fill(prompt); diff --git a/tests/e2e/ui/tests/internal-user/internalUser.spec.ts b/tests/e2e/ui/tests/internal-user/internalUser.spec.ts index f392c5104da..d34bc034453 100644 --- a/tests/e2e/ui/tests/internal-user/internalUser.spec.ts +++ b/tests/e2e/ui/tests/internal-user/internalUser.spec.ts @@ -9,7 +9,7 @@ import { import { Page } from "../../fixtures/pages"; import { navigateToPage, clickTeamId } from "../../helpers/navigation"; import { CHAT_MODEL_A, MOCK_RESPONSE_TEXT, masterKey } from "../../helpers/traffic"; -import { keySourceSelect, onlyVisible, openPlayground, selectModel, sendMessage } from "../../helpers/playground"; +import { openPlayground, selectModel, sendMessage, enterVirtualKey } from "../../helpers/playground"; test.describe("Internal User", () => { test.use({ storageState: INTERNAL_USER_STORAGE_PATH }); @@ -73,12 +73,7 @@ test.describe("Internal User", () => { try { await openPlayground(page); - await keySourceSelect(page).click(); - await onlyVisible(page.getByRole("option", { name: "Virtual Key" })).click({ timeout: 15_000 }); - - const keyInput = onlyVisible(page.getByPlaceholder("Enter custom Virtual Key")); - await expect(keyInput).toBeVisible({ timeout: 10_000 }); - await keyInput.fill(apiKey); + await enterVirtualKey(page, apiKey); await selectModel(page, CHAT_MODEL_A); await sendMessage(page, `internal user team key ping ${keyName}`); diff --git a/tests/e2e/ui/tests/playground/playground.spec.ts b/tests/e2e/ui/tests/playground/playground.spec.ts index 9a9ecd539b3..205cc57ae46 100644 --- a/tests/e2e/ui/tests/playground/playground.spec.ts +++ b/tests/e2e/ui/tests/playground/playground.spec.ts @@ -1,7 +1,14 @@ import { test, expect } from "@playwright/test"; import { ADMIN_STORAGE_PATH } from "../../constants"; import { CHAT_MODEL_A, CHAT_MODEL_B, MOCK_RESPONSE_TEXT, createVirtualKey } from "../../helpers/traffic"; -import { keySourceSelect, onlyVisible, openPlayground, selectModel, sendMessage } from "../../helpers/playground"; +import { + keySourceSelect, + openPlayground, + openRequestSettings, + selectModel, + sendMessage, + enterVirtualKey, +} from "../../helpers/playground"; /** * The one flow that exercises the dashboard's own LLM call path rather than an admin CRUD endpoint, @@ -14,8 +21,9 @@ test.describe("Playground", () => { test(`chats with ${model} using the current UI session`, async ({ page }) => { await openPlayground(page); - // "Current UI Session" is the default: the logged-in admin's key, nothing pasted. - await expect(keySourceSelect(page)).toContainText("Current UI Session"); + await openRequestSettings(page); + await expect(keySourceSelect(page)).toContainText("Current UI session"); + await page.keyboard.press("Escape"); await selectModel(page, model); const prompt = `playground ping for ${model}`; @@ -34,14 +42,7 @@ test.describe("Playground", () => { await openPlayground(page); - // Switch the source to "Virtual Key" and paste the key we just minted. - await keySourceSelect(page).click(); - await onlyVisible(page.getByRole("option", { name: "Virtual Key" })).click({ timeout: 15_000 }); - - const keyInput = onlyVisible(page.getByPlaceholder("Enter custom Virtual Key")); - await expect(keyInput).toBeVisible({ timeout: 10_000 }); - await keyInput.fill(key); - + await enterVirtualKey(page, key); await selectModel(page, CHAT_MODEL_A); await sendMessage(page, "playground ping via virtual key"); diff --git a/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts b/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts index f3c031f0172..ff41778b2c9 100644 --- a/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts +++ b/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts @@ -10,7 +10,7 @@ import { Page } from "../../fixtures/pages"; import { navigateToPage, dismissFeedbackPopup, clickTeamId } from "../../helpers/navigation"; import { captureRequestBody, readBack } from "../../helpers/roundTrip"; import { CHAT_MODEL_A, masterKey } from "../../helpers/traffic"; -import { keySourceSelect, modelSelect, onlyVisible, openPlayground } from "../../helpers/playground"; +import { onlyVisible, openPlayground, searchModels, enterVirtualKey } from "../../helpers/playground"; /** * Every identifier a roster is addressable by. Which of user_id / user_email is populated depends on @@ -228,21 +228,14 @@ test.describe("Team Admin", () => { .toBe(true); await openPlayground(page); - await keySourceSelect(page).click(); - await onlyVisible(page.getByRole("option", { name: "Virtual Key" })).click({ timeout: 15_000 }); + await enterVirtualKey(page, teamKey); - const keyInput = onlyVisible(page.getByPlaceholder("Enter custom Virtual Key")); - await expect(keyInput).toBeVisible({ timeout: 10_000 }); - await keyInput.fill(teamKey); - - const select = modelSelect(page); - await select.click(); - await select.fill(teamModelName); + await searchModels(page, teamModelName); await expect(onlyVisible(page.getByRole("option", { name: teamModelName }))).toBeVisible({ timeout: 15_000, }); - await select.fill(CHAT_MODEL_A); + await onlyVisible(page.getByPlaceholder("Search models")).fill(CHAT_MODEL_A); await expect(onlyVisible(page.getByRole("option", { name: CHAT_MODEL_A }))).toBeVisible({ timeout: 15_000, }); diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index b1b6cef7320..9eead41a102 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -2090,11 +2090,26 @@ "count": 1 } }, + "src/components/ui/attachment.tsx": { + "local/no-ad-hoc-z-index": { + "count": 2 + } + }, + "src/components/ui/message-scroller.tsx": { + "local/no-noop-hover-variant": { + "count": 1 + } + }, "src/components/ui/slider.tsx": { "no-nested-ternary": { "count": 1 } }, + "src/components/ui/toggle-group.tsx": { + "local/no-ad-hoc-z-index": { + "count": 2 + } + }, "src/components/update_model_credentials_modal.tsx": { "local/filename-pascal-case": { "count": 1 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/ByokConnect.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/ByokConnect.tsx new file mode 100644 index 00000000000..f94c43b720a --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/ByokConnect.tsx @@ -0,0 +1,152 @@ +"use client"; + +import { useId, useState } from "react"; +import { Check, ExternalLink, KeyRound } from "lucide-react"; +import { PasswordInput } from "@/components/shared/PasswordInput"; +import { Button } from "@/components/ui/button"; +import { InputGroupButton } from "@/components/ui/input-group"; +import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover"; +import { Field, FieldDescription, FieldGroup, FieldLabel } from "@/components/ui/field"; +import { Spinner } from "@/components/ui/spinner"; +import { Switch } from "@/components/ui/switch"; +import type { MCPServer } from "@/components/mcp_tools/types"; +import { fetchClient } from "@/lib/http/api"; +import { ApiError } from "@/lib/http/client"; +import { toast } from "@/lib/toast"; + +const saveErrorMessage = (error: unknown): string => { + if (error instanceof ApiError) { + const detail = (error.body as { detail?: { error?: string } } | null)?.detail?.error; + if (detail) return detail; + } + return error instanceof Error && error.message ? error.message : "Failed to connect"; +}; + +export function ByokConnect({ + server, + name, + onConnected, + onCancel, +}: { + server: MCPServer; + name: string; + onConnected: () => void; + onCancel?: () => void; +}) { + const [credential, setCredential] = useState(""); + const [save, setSave] = useState(true); + const [isSaving, setIsSaving] = useState(false); + const keyId = useId(); + const saveId = useId(); + + const connect = async () => { + if (!credential.trim()) { + toast.error("Enter your API key"); + return; + } + setIsSaving(true); + try { + await fetchClient.POST("/v1/mcp/server/{server_id}/user-credential", { + params: { path: { server_id: server.server_id } }, + body: { credential: credential.trim(), save }, + }); + toast.success(`Connected to ${name}`); + setCredential(""); + onConnected(); + } catch (error) { + toast.error(saveErrorMessage(error)); + } finally { + setIsSaving(false); + } + }; + + return ( + +
+

+ + {server.has_user_credential ? `Replace your ${name} API key` : `${name} needs your API key`} +

+

+ {server.has_user_credential + ? `The new key replaces the one the gateway uses for ${name}.` + : `The gateway calls ${name} with your own key. Connect it to list and use its tools.`} +

+
+ {server.byok_description && server.byok_description.length > 0 && ( +
    + {server.byok_description.map((item) => ( +
  • + + {item} +
    • + ))} +
+ )} + + {name} API key + setCredential(event.target.value)} + placeholder="Paste your API key" + /> + {server.byok_api_key_help_url && ( + + + Where to find this key + + + + )} + + + + Save the key for next time + + + +
+ + {onCancel && ( + + )} +
+ + ); +} + +export function ByokKeyButton({ + server, + name, + onConnected, +}: { + server: MCPServer; + name: string; + onConnected: () => void; +}) { + const [open, setOpen] = useState(false); + return ( + + }> + + {server.has_user_credential ? "Replace key" : "Connect key"} + + + { + setOpen(false); + onConnected(); + }} + /> + + + ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/Composer.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/Composer.test.tsx new file mode 100644 index 00000000000..c62b839e5b3 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/Composer.test.tsx @@ -0,0 +1,61 @@ +import { createRef } from "react"; +import { createPortal } from "react-dom"; +import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { describe, expect, it, vi } from "vitest"; +import { EndpointType } from "@/components/chat_ui/mode_endpoint_mapping"; +import type { MCPToolArgumentsFormRef } from "@/components/mcp_tools/MCPToolArgumentsForm"; +import { Composer, type ChecklistProps, type DraftState } from "./Composer"; +import { capabilitiesFor } from "./endpointCapabilities"; + +const checklist: ChecklistProps = { groups: [], value: [], onValueChange: vi.fn(), isLoading: false }; + +const draft: DraftState = { + text: "", + setText: vi.fn(), + files: [], + attach: vi.fn(), + removeFile: vi.fn(), + isLoading: false, + send: vi.fn().mockResolvedValue(undefined), + cancel: vi.fn(), +}; + +const renderComposer = () => + render( + } + draft={draft} + directTool={null} + toolArgsForm={createRef()} + mcpControl={null} + vectorStores={checklist} + guardrails={checklist} + codeInterpreter={{ enabled: false, onToggle: vi.fn() }} + settingsControl={createPortal(, document.body)} + sendDisabled + />, + ); + +describe("Composer", () => { + it("keeps focus in a popover field opened from the composer toolbar", async () => { + const user = userEvent.setup(); + renderComposer(); + + await user.click(screen.getByLabelText("Request tags")); + + expect(screen.getByLabelText("Request tags")).toHaveFocus(); + }); + + it("focuses the message box when empty toolbar space is clicked", async () => { + const user = userEvent.setup(); + const { container } = renderComposer(); + const toolbar = container.querySelectorAll('[data-slot="input-group-addon"]'); + + await user.click(toolbar[toolbar.length - 1]); + + expect(screen.getByLabelText("Message")).toHaveFocus(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/Composer.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/Composer.tsx new file mode 100644 index 00000000000..a9fbb777393 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/Composer.tsx @@ -0,0 +1,196 @@ +"use client"; + +import { useRef, type DragEvent, type KeyboardEvent, type MouseEvent, type ReactNode, type RefObject } from "react"; +import { ArrowUp, Code2, Database, Paperclip, Shield, Square } from "lucide-react"; +import { EndpointType } from "@/components/chat_ui/mode_endpoint_mapping"; +import MCPToolArgumentsForm, { type MCPToolArgumentsFormRef } from "@/components/mcp_tools/MCPToolArgumentsForm"; +import type { MCPTool } from "@/components/mcp_tools/types"; +import { InputGroup, InputGroupAddon, InputGroupButton, InputGroupTextarea } from "@/components/ui/input-group"; +import { Toggle } from "@/components/ui/toggle"; +import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; +import { DraftAttachments } from "./ComposerControls"; +import type { DraftFile } from "./composerFiles"; +import { ACCEPT_FOR, ATTACH_LABEL, type Capabilities } from "./endpointCapabilities"; +import { ChecklistMenu, type OptionGroup } from "./Pickers"; + +export type ChecklistProps = Readonly<{ + groups: readonly OptionGroup[]; + value: readonly string[]; + onValueChange: (value: readonly string[]) => void; + isLoading: boolean; +}>; + +export type DraftState = Readonly<{ + text: string; + setText: (text: string) => void; + files: readonly DraftFile[]; + attach: (files: readonly File[]) => void; + removeFile: (index: number) => void; + isLoading: boolean; + send: () => Promise; + cancel: () => void; +}>; + +function AttachButton({ + capabilities, + attach, +}: { + capabilities: Capabilities; + attach: (files: readonly File[]) => void; +}) { + const input = useRef(null); + if (!capabilities.attach) return null; + const label = ATTACH_LABEL[capabilities.attach]; + return ( + <> + + input.current?.click()} />} + > + + + {label} + + { + attach(Array.from(event.target.files ?? [])); + event.target.value = ""; + }} + /> + + ); +} + +function SendButton({ draft, disabled }: { draft: DraftState; disabled: boolean }) { + if (draft.isLoading) + return ( + + + + ); + return ( + void draft.send()} + > + + + ); +} + +export function Composer({ + capabilities, + endpoint, + targetRow, + draft, + directTool, + toolArgsForm, + mcpControl, + vectorStores, + guardrails, + codeInterpreter, + settingsControl, + sendDisabled, +}: { + capabilities: Capabilities; + endpoint: EndpointType; + targetRow: ReactNode; + draft: DraftState; + directTool: MCPTool | null; + toolArgsForm: RefObject; + mcpControl: ReactNode; + vectorStores: ChecklistProps; + guardrails: ChecklistProps; + codeInterpreter: Readonly<{ enabled: boolean; onToggle: (enabled: boolean) => void }>; + settingsControl: ReactNode; + sendDisabled: boolean; +}) { + const isDirectCall = endpoint === EndpointType.MCP; + const isRealtime = endpoint === EndpointType.REALTIME; + const message = useRef(null); + const focusMessage = (event: MouseEvent) => { + const target = event.target as HTMLElement; + if (!event.currentTarget.contains(target) || target.closest("button, input, textarea")) return; + message.current?.focus(); + }; + const handleDrop = (event: DragEvent) => { + if (!capabilities.attach) return; + event.preventDefault(); + draft.attach(Array.from(event.dataTransfer.files)); + }; + const handleKeyDown = (event: KeyboardEvent) => { + if (event.key !== "Enter" || event.shiftKey || event.nativeEvent.isComposing) return; + event.preventDefault(); + if (!sendDisabled) void draft.send(); + }; + const directCallBody = ( +
+ {directTool ? ( + + ) : ( +

Pick a server and a tool above to fill in its arguments

+ )} +
+ ); + const textBody = ( + draft.setText(event.target.value)} + onKeyDown={handleKeyDown} + placeholder={capabilities.placeholder} + rows={2} + aria-label="Message" + className="max-h-60" + /> + ); + return ( + capabilities.attach && event.preventDefault()} onDrop={handleDrop}> + + {targetRow} + + {draft.files.length > 0 && ( + + + + )} + {!isRealtime && (isDirectCall ? directCallBody : textBody)} + + {capabilities.mcp && mcpControl} + {capabilities.vectorStores && ( + } emptyText="No vector stores yet" {...vectorStores} /> + )} + {capabilities.guardrails && ( + } emptyText="No guardrails yet" {...guardrails} /> + )} + {capabilities.codeInterpreter && ( + + + Code interpreter + + )} + + {settingsControl} + + {!isRealtime && } + + + + ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/ComposerControls.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/ComposerControls.tsx new file mode 100644 index 00000000000..e3e52be3b7b --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/ComposerControls.tsx @@ -0,0 +1,120 @@ +"use client"; + +import { ChevronDown, FileText, Music, X } from "lucide-react"; +import { EndpointType } from "@/components/chat_ui/mode_endpoint_mapping"; +import { + Attachment, + AttachmentAction, + AttachmentActions, + AttachmentContent, + AttachmentDescription, + AttachmentGroup, + AttachmentMedia, + AttachmentTitle, +} from "@/components/ui/attachment"; +import { + DropdownMenu, + DropdownMenuContent, + DropdownMenuLabel, + DropdownMenuRadioGroup, + DropdownMenuRadioItem, + DropdownMenuTrigger, +} from "@/components/ui/dropdown-menu"; +import { InputGroupButton } from "@/components/ui/input-group"; +import { ENDPOINT_OPTIONS, OPEN_AI_VOICE_SELECT_OPTIONS, type OpenAIVoice } from "../chat_ui/chatConstants"; +import { formatBytes, type DraftFile } from "./composerFiles"; + +const shortEndpoint = (label: string): string => label.replace(/^\/v1(beta)?\//, ""); + +export function EndpointMenu({ + value, + onValueChange, +}: { + value: EndpointType; + onValueChange: (value: EndpointType) => void; +}) { + const current = ENDPOINT_OPTIONS.find((option) => option.value === value); + return ( + + }> + {shortEndpoint(current?.label ?? value)} + + + + onValueChange(next as EndpointType)}> + Endpoint + {ENDPOINT_OPTIONS.map((option) => ( + + {option.label} + + ))} + + + + ); +} + +export function VoiceMenu({ + value, + onValueChange, +}: { + value: OpenAIVoice; + onValueChange: (value: OpenAIVoice) => void; +}) { + const current = OPEN_AI_VOICE_SELECT_OPTIONS.find((option) => option.value === value); + return ( + + }> + Voice: {current?.label.split(" - ")[0] ?? value} + + + + onValueChange(next as OpenAIVoice)}> + Voice + {OPEN_AI_VOICE_SELECT_OPTIONS.map((option) => ( + + {option.label} + + ))} + + + + ); +} + +function FileMedia({ draft }: { draft: DraftFile }) { + if (draft.previewUrl) + return ( + + + + ); + return {draft.file.type.startsWith("audio/") ? : }; +} + +export function DraftAttachments({ + files, + onRemove, +}: { + files: readonly DraftFile[]; + onRemove: (index: number) => void; +}) { + return ( + + {files.map((draft, index) => ( + + + + {draft.file.name} + {formatBytes(draft.file.size)} + + + onRemove(index)}> + + + + + ))} + + ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/Conversation.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/Conversation.tsx new file mode 100644 index 00000000000..31eb395408a --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/Conversation.tsx @@ -0,0 +1,388 @@ +"use client"; + +import type { ReactNode } from "react"; +import { + ArrowDownToLine, + ArrowUpFromLine, + Brain, + Clock, + Coins, + DatabaseZap, + FileText, + Music, + PlayCircle, + Repeat, + ScanSearch, + Timer, + Wrench, +} from "lucide-react"; +import { EndpointType } from "@/components/chat_ui/mode_endpoint_mapping"; +import MCPEventsDisplay from "@/components/chat_ui/MCPEventsDisplay"; +import ReasoningContent from "@/components/chat_ui/ReasoningContent"; +import type { MessageAttachment, MessageType } from "@/components/chat_ui/types"; +import type { CodeInterpreterResult } from "@/components/llm_calls/code_interpreter_handler"; +import type { MCPEvent } from "@/components/mcp_tools/types"; +import { Alert, AlertDescription, AlertTitle } from "@/components/ui/alert"; +import { + Attachment, + AttachmentContent, + AttachmentDescription, + AttachmentGroup, + AttachmentMedia, + AttachmentTitle, +} from "@/components/ui/attachment"; +import { Badge } from "@/components/ui/badge"; +import { Bubble, BubbleContent } from "@/components/ui/Bubble"; +import { Button } from "@/components/ui/button"; +import { Empty, EmptyContent, EmptyDescription, EmptyHeader, EmptyMedia, EmptyTitle } from "@/components/ui/empty"; +import { Marker, MarkerContent, MarkerIcon } from "@/components/ui/marker"; +import { Message, MessageContent, MessageFooter, MessageHeader } from "@/components/ui/message"; +import { + MessageScroller, + MessageScrollerButton, + MessageScrollerContent, + MessageScrollerItem, + MessageScrollerProvider, + MessageScrollerViewport, +} from "@/components/ui/message-scroller"; +import { Spinner } from "@/components/ui/spinner"; +import { Toggle } from "@/components/ui/toggle"; +import A2AMetrics from "../chat_ui/A2AMetrics"; +import CodeInterpreterOutput from "../chat_ui/CodeInterpreterOutput"; +import { SearchResultsDisplay } from "../chat_ui/SearchResultsDisplay"; +import { formatBytes } from "./composerFiles"; +import { CopyIconButton, FallbackBadge, LongValue } from "./ModelName"; +import { ReplyBody } from "./ReplyContent"; + +export const formatSeconds = (ms: number): string => `${(ms / 1000).toFixed(2)}s`; + +export const formatCost = (cost: number): string => + `$${cost < 0.0001 && cost > 0 ? cost.toPrecision(2) : cost.toFixed(4)}`; + +const MCP_EVENT_ENDPOINTS: readonly EndpointType[] = [EndpointType.CHAT, EndpointType.RESPONSES]; + +const textOf = (message: MessageType): string => (typeof message.content === "string" ? message.content : ""); + +const fellBack = (message: MessageType): boolean => (message.trace?.attemptedFallbacks ?? 0) > 0; + +export const servedModel = (message: MessageType): string | undefined => + message.trace?.modelGroup ?? (message.isError ? undefined : message.model ?? message.trace?.requestedModel); + +const replyModel = (message: MessageType): Readonly<{ value: string; detail: string }> | null => { + const served = servedModel(message); + if (served) return { value: served, detail: fellBack(message) ? "Served by, after a fallback" : "Served by" }; + const requested = message.trace?.requestedModel ?? message.model; + return requested ? { value: requested, detail: "Requested model" } : null; +}; + +function Metric({ icon, value, label }: { icon: ReactNode; value: string; label?: string }) { + return ( + + {icon} + {value} + {label} + + ); +} + +function ReplyMetrics({ message }: { message: MessageType }) { + const usage = message.usage; + const cost = usage?.cost ?? message.trace?.cost; + const cached = usage?.servedFromResponseCache || message.trace?.cacheHit; + return ( +
+ {message.timeToFirstToken !== undefined && ( + } value={formatSeconds(message.timeToFirstToken)} label="to first token" /> + )} + {message.totalLatency !== undefined && ( + } value={formatSeconds(message.totalLatency)} label="total" /> + )} + {usage?.promptTokens !== undefined && ( + } value={String(usage.promptTokens)} label="in" /> + )} + {usage?.completionTokens !== undefined && ( + } value={String(usage.completionTokens)} label="out" /> + )} + {usage?.reasoningTokens ? ( + } value={String(usage.reasoningTokens)} label="reasoning" /> + ) : null} + {usage?.cacheReadTokens ? ( + } value={String(usage.cacheReadTokens)} label="cached" /> + ) : null} + {cost !== undefined && } value={formatCost(cost)} />} + {message.toolName && } value={message.toolName} />} + {cached && Cached response} +
+ ); +} + +function AttachmentMediaFor({ attachment }: { attachment: MessageAttachment }) { + if (attachment.previewUrl) + return ( + + + + ); + return {attachment.kind === "audio" ? : }; +} + +function UserTurn({ message }: { message: MessageType }) { + const text = textOf(message); + return ( + + + {message.attachments && message.attachments.length > 0 && ( + + {message.attachments.map((attachment, index) => ( + + + + {attachment.name} + {formatBytes(attachment.size)} + + + ))} + + )} + {message.mcpToolCall && ( + + + + + Called {message.mcpToolCall.tool} + + 
+                {JSON.stringify(message.mcpToolCall.arguments, null, 2)}
+
+ + + )} + {text && ( + + {text} + + )} + + + ); +} + +export type LastReplyExtras = Readonly<{ + endpoint: EndpointType; + mcpEvents: readonly MCPEvent[]; + codeInterpreter: CodeInterpreterResult | null; + accessToken: string; +}>; + +function AssistantReply({ + message, + isLast, + extras, + inspected, + onInspect, +}: { + message: MessageType; + isLast: boolean; + extras: LastReplyExtras; + inspected: boolean; + onInspect: () => void; +}) { + const model = replyModel(message); + const showMcpEvents = isLast && extras.mcpEvents.length > 0 && MCP_EVENT_ENDPOINTS.includes(extras.endpoint); + const showCode = isLast && extras.codeInterpreter && extras.endpoint === EndpointType.RESPONSES; + return ( + + + {model && ( + + + {fellBack(message) && } + + )} + {message.reasoningContent && } + {showMcpEvents && } + {message.searchResults && } + {showCode && extras.codeInterpreter && ( + + )} + {message.isError ? ( + + The request failed + {textOf(message)} + + ) : ( + + )} + +
+ {message.a2aMetadata ? ( + + ) : ( + + )} +
+ + {!message.isImage && !message.isAudio && } + {message.trace && ( + + + Inspect + + )} + + + + + ); +} + +function WaitingTurn({ target }: { target: string | undefined }) { + return ( + + +
+ + {target ? `Waiting for ${target}` : "Waiting for a reply"} +
+ + + ); +} + +function SwitchMarker({ model }: { model: string }) { + return ( + + + + + Switched to {model} + + ); +} + +function EmptyConversation({ + suggestions, + onPick, +}: { + suggestions: readonly string[]; + onPick: (text: string) => void; +}) { + return ( + + + + + + Send a test request + + Pick a model and anything the request should go through, then inspect what the gateway did with it + + + {suggestions.length > 0 && ( + + {suggestions.map((suggestion) => ( + + ))} + + )} + + ); +} + +const switchedModel = (messages: readonly MessageType[], index: number): string | null => { + const message = messages[index]; + if (message?.role !== "user" || !message.model) return null; + const previous = messages.slice(0, index).findLast((candidate) => candidate.role === "user"); + return previous?.model && previous.model !== message.model ? message.model : null; +}; + +export function Conversation({ + messages, + isLoading, + waitingFor, + extras, + inspectedIndex, + onInspect, + suggestions, + onPickSuggestion, +}: { + messages: readonly MessageType[]; + isLoading: boolean; + waitingFor: string | undefined; + extras: LastReplyExtras; + inspectedIndex: number | null; + onInspect: (index: number) => void; + suggestions: readonly string[]; + onPickSuggestion: (text: string) => void; +}) { + if (messages.length === 0) return ; + const lastIndex = messages.length - 1; + const waiting = isLoading && messages[lastIndex]?.role === "user"; + return ( + + + + + {messages.flatMap((message, index) => { + const switched = switchedModel(messages, index); + const turn = ( + + {message.role === "user" ? ( + + ) : ( + onInspect(index)} + /> + )} + + ); + return switched + ? [ + + + , + turn, + ] + : [turn]; + })} + {waiting && ( + + + + )} + + + + + + ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/Inspector.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/Inspector.tsx new file mode 100644 index 00000000000..3a9919ce38f --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/Inspector.tsx @@ -0,0 +1,292 @@ +"use client"; + +import { useState, type ReactNode } from "react"; +import { + ArrowDownToLine, + ArrowUpFromLine, + Brain, + Clock, + Coins, + DatabaseZap, + Gauge, + RotateCw, + Server, + Shuffle, + Timer, +} from "lucide-react"; +import type { MessageType } from "@/components/chat_ui/types"; +import { Accordion, AccordionContent, AccordionItem, AccordionTrigger } from "@/components/ui/accordion"; +import { Badge } from "@/components/ui/badge"; +import { Empty, EmptyDescription, EmptyHeader, EmptyTitle } from "@/components/ui/empty"; +import { FieldDescription } from "@/components/ui/field"; +import { Item, ItemActions, ItemContent, ItemDescription, ItemGroup, ItemTitle } from "@/components/ui/item"; +import { ScrollArea, ScrollBar } from "@/components/ui/scroll-area"; +import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; +import { ToggleGroup, ToggleGroupItem } from "@/components/ui/toggle-group"; +import { formatCost, formatSeconds, servedModel } from "./Conversation"; +import { CopyIconButton, FallbackBadge, LongValue } from "./ModelName"; + +export type CodeSdk = "openai" | "azure"; + +function TraceRow({ title, description, children }: { title: string; description?: string; children?: ReactNode }) { + return ( + + + {title} + {description && {description}} + + {children && {children}} + + ); +} + +type StatValue = Readonly<{ label: string; icon: ReactNode; value: string | undefined }>; + +function StatGrid({ stats }: { stats: readonly StatValue[] }) { + const shown = stats.filter((stat): stat is StatValue & { value: string } => stat.value !== undefined); + if (shown.length === 0) return null; + return ( + + {shown.map((stat) => ( + + + {stat.icon} + {stat.label} + + {stat.value} + + ))} + + ); +} + +function TraceSection({ value, title, children }: { value: string; title: string; children: ReactNode }) { + return ( + + {title} + + {children} + + + ); +} + +function Muted({ children }: { children: ReactNode }) { + return

{children}

; +} + +const count = (value: number | undefined): string | undefined => (value === undefined ? undefined : String(value)); +const millis = (value: number | undefined): string | undefined => + value === undefined ? undefined : `${Math.round(value)}ms`; + +function Routing({ message }: { message: MessageType }) { + const trace = message.trace; + const served = servedModel(message); + const fellBack = (trace?.attemptedFallbacks ?? 0) > 0; + const reported = Boolean(trace?.callId || trace?.modelGroup || trace?.providerModel); + return ( + + {trace?.requestedModel && ( + + + + )} + {served && ( + + {fellBack && } + + + )} + {trace?.providerModel && ( + + + + )} + {trace?.apiBase && ( + + + + )} + {trace?.callId && ( + + {trace.callId} + + + )} + {trace?.responseId && ( + + {trace.responseId} + + + )} + , value: count(trace?.attemptedFallbacks) }, + { label: "Retries", icon: , value: count(trace?.attemptedRetries) }, + { label: "Overhead", icon: , value: millis(trace?.overheadMs) }, + ]} + /> + {!reported && ( + + The gateway did not share routing details with this page. Open the dashboard from the proxy to see them + + )} + + ); +} + +const noneRan = (policies: readonly string[] | undefined): string | null => { + if (policies === undefined) return "No guardrails ran on this request"; + return policies.length === 0 ? "No guardrails or policies ran on this request" : null; +}; + +function Guardrails({ message }: { message: MessageType }) { + const guardrails = message.trace?.appliedGuardrails; + const policies = message.trace?.appliedPolicies; + if (guardrails === undefined) { + return ( + + This endpoint does not report which guardrails or policies ran + + ); + } + const emptyNote = guardrails.length === 0 ? noneRan(policies) : null; + return ( + + {guardrails.map((name) => ( + + Ran + + ))} + {policies?.map((name) => ( + + Ran + + ))} + {emptyNote && {emptyNote}} + {policies === undefined && Policies are only reported when the dashboard is opened from the proxy} + + ); +} + +function SpeedAndCost({ message }: { message: MessageType }) { + const usage = message.usage; + const cost = usage?.cost ?? message.trace?.cost; + return ( + + , + value: message.timeToFirstToken === undefined ? undefined : formatSeconds(message.timeToFirstToken), + }, + { + label: "Total", + icon: , + value: message.totalLatency === undefined ? undefined : formatSeconds(message.totalLatency), + }, + { label: "Cost", icon: , value: cost === undefined ? undefined : formatCost(cost) }, + { label: "Tokens in", icon: , value: count(usage?.promptTokens) }, + { label: "Tokens out", icon: , value: count(usage?.completionTokens) }, + { label: "Cached", icon: , value: count(usage?.cacheReadTokens) }, + { label: "Reasoning", icon: , value: count(usage?.reasoningTokens) }, + { label: "Gateway time", icon: , value: millis(message.trace?.responseMs) }, + ]} + /> + + ); +} + +function NothingYet() { + return ( + + + Nothing to inspect yet + + Each reply shows here with the deployment that served it, fallbacks, guardrails, timing and cost + + + + ); +} + +function Trace({ message }: { message: MessageType | null }) { + if (!message) return ; + return ( + + + + + + ); +} + +function CodeTab({ code }: { code: (sdk: CodeSdk) => string }) { + const [sdk, setSdk] = useState("openai"); + const snippet = code(sdk); + return ( +
+
+ value[0] && setSdk(value[0] as CodeSdk)} + variant="outline" + size="sm" + spacing={0} + aria-label="SDK" + > + + OpenAI SDK + + + Azure SDK + + + + + +
+ + Send this request from your own app. Replace YOUR_LITELLM_API_KEY with a virtual key + + + 
{snippet}
+ + +
+ ); +} + +export function InspectorPanel({ message, code }: { message: MessageType | null; code: (sdk: CodeSdk) => string }) { + return ( + +
+ + Trace + Code + +
+ + + + + + + + + + ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/McpDialog.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/McpDialog.tsx new file mode 100644 index 00000000000..e5e0828825d --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/McpDialog.tsx @@ -0,0 +1,391 @@ +"use client"; + +import { useId, useState, type ReactNode } from "react"; +import { Blocks, KeyRound, Search } from "lucide-react"; +import { Badge } from "@/components/ui/badge"; +import { Button } from "@/components/ui/button"; +import { Checkbox } from "@/components/ui/checkbox"; +import { + Dialog, + DialogClose, + DialogContent, + DialogDescription, + DialogFooter, + DialogHeader, + DialogTitle, + DialogTrigger, +} from "@/components/ui/dialog"; +import { Empty, EmptyContent, EmptyDescription, EmptyHeader, EmptyTitle } from "@/components/ui/empty"; +import { Field, FieldLabel, FieldLegend, FieldSet } from "@/components/ui/field"; +import { InputGroup, InputGroupAddon, InputGroupButton, InputGroupInput } from "@/components/ui/input-group"; +import { ScrollArea } from "@/components/ui/scroll-area"; +import { Spinner } from "@/components/ui/spinner"; +import type { MCPServer, MCPToolset } from "@/components/mcp_tools/types"; +import { uiHref } from "@/utils/uiHref"; +import { ByokConnect } from "./ByokConnect"; +import { + ALL_MCP_SERVERS, + TOOLSET_PREFIX, + allowedTools, + toggleMcpItem, + toggleToolLimit, + withoutStaleLimits, + type ToolLimits, +} from "./playgroundSettings"; +import { useServerTools } from "./usePlaygroundCatalog"; + +type Row = Readonly<{ value: string; label: string }>; + +export const serverName = (server: MCPServer): string => server.alias || server.server_name || server.server_id; + +const matches = (value: string, query: string): boolean => value.toLowerCase().includes(query.trim().toLowerCase()); + +const needsKey = (server: MCPServer | undefined): boolean => Boolean(server?.is_byok && !server.has_user_credential); + +function SearchField({ + value, + onChange, + placeholder, + autoFocus, +}: { + value: string; + onChange: (value: string) => void; + placeholder: string; + autoFocus: boolean; +}) { + return ( + + + + + onChange(event.target.value)} + placeholder={placeholder} + aria-label={placeholder} + autoFocus={autoFocus} + /> + + ); +} + +function MasterRow({ + row, + checked, + onCheckedChange, + active, + onSelect, + summary, +}: { + row: Row; + checked: boolean; + onCheckedChange: (checked: boolean) => void; + active: boolean; + onSelect: () => void; + summary?: string; +}) { + return ( +
+ + +
+ ); +} + +function ToolRow({ + label, + checked, + disabled, + onCheckedChange, +}: { + label: string; + checked: boolean; + disabled: boolean; + onCheckedChange: (checked: boolean) => void; +}) { + const id = useId(); + return ( + + + + {label} + + + ); +} + +function Note({ title, children, action }: { title: string; children: ReactNode; action?: ReactNode }) { + return ( + + + {title} + {children} + + {action && {action}} + + ); +} + +function RowGroup({ title, rows, render }: { title: string; rows: readonly Row[]; render: (row: Row) => ReactNode }) { + if (rows.length === 0) return null; + return ( +
+ + {title} + + {rows.map(render)} +
+ ); +} + +export function McpDialog({ + requestKey, + servers, + toolsets, + isLoading, + selection, + toolLimits, + onChange, + onCredentialSaved, + defaultOpen = false, +}: { + requestKey: string | null; + servers: readonly MCPServer[]; + toolsets: readonly MCPToolset[]; + isLoading: boolean; + selection: readonly string[]; + toolLimits: ToolLimits; + onChange: (selection: readonly string[], limits: ToolLimits) => void; + onCredentialSaved: () => void; + defaultOpen?: boolean; +}) { + const [focused, setFocused] = useState(selection.find((item) => item !== ALL_MCP_SERVERS) ?? ALL_MCP_SERVERS); + const [query, setQuery] = useState(""); + const [toolQuery, setToolQuery] = useState(""); + const [replacingKey, setReplacingKey] = useState(null); + const serverById = new Map(servers.map((server) => [server.server_id, server])); + const selectedServers = selection.filter((item) => serverById.has(item)); + const toolsByServer = useServerTools( + requestKey, + [...new Set([...selectedServers, ...(serverById.has(focused) ? [focused] : [])])].filter( + (id) => !needsKey(serverById.get(id)), + ), + ); + const toolNames = (serverId: string): readonly string[] => + (toolsByServer[serverId]?.tools ?? []).map((tool) => tool.name); + + const serverRows: Row[] = servers.map((server) => ({ value: server.server_id, label: serverName(server) })); + const toolsetRows: Row[] = toolsets.map((toolset) => ({ + value: `${TOOLSET_PREFIX}${toolset.toolset_id}`, + label: toolset.toolset_name, + })); + const visible = (rows: readonly Row[]) => rows.filter((row) => matches(row.label, query)); + + const toggle = (item: string, on: boolean) => { + const next = toggleMcpItem(selection, item, on); + onChange(next, withoutStaleLimits(toolLimits, next)); + if (on) setFocused(item); + }; + + const summary = (row: Row): string | undefined => { + if (!selection.includes(row.value)) return undefined; + const server = serverById.get(row.value); + if (!server) return undefined; + if (needsKey(server)) return "Needs your key"; + const tools = toolNames(row.value); + if (tools.length === 0) return undefined; + const allowed = allowedTools(toolLimits, row.value, tools).length; + return allowed === tools.length ? `All ${tools.length} tools` : `${allowed} of ${tools.length} tools`; + }; + + const focus = (item: string) => { + setFocused(item); + setToolQuery(""); + setReplacingKey(null); + }; + + const masterRow = (row: Row) => ( + toggle(row.value, on)} + active={row.value === focused} + onSelect={() => focus(row.value)} + summary={summary(row)} + /> + ); + + const detail = (): ReactNode => { + if (focused === ALL_MCP_SERVERS) + return Every server this key can reach, with all of its tools; + const toolset = toolsets.find((item) => `${TOOLSET_PREFIX}${item.toolset_id}` === focused); + if (toolset) + return ( + +
+

+ {toolset.description || "A toolset's tools are fixed by whoever created it"} +

+
+ {toolset.tools.map((tool) => ( + + {tool.tool_name} + + ))} +
+
+ + ); + const server = serverById.get(focused); + if (!server) return Pick a server or toolset on the left; + const name = serverName(server); + if (needsKey(server) || replacingKey === focused) + return ( + { + setReplacingKey(null); + onCredentialSaved(); + }} + onCancel={needsKey(server) ? undefined : () => setReplacingKey(null)} + /> + ); + const replaceKey = server.is_byok ? ( + + ) : undefined; + if (!selection.includes(focused)) + return ( + + Check this server to choose which of its tools the model can call + + ); + const state = toolsByServer[focused]; + if (!state || state.isLoading) + return ( +
+ +
+ ); + if (state.failed) + return ( + + Couldn't load this server's tools + + ); + const tools = toolNames(focused); + if (tools.length === 0) return This server has no tools; + const allowed = allowedTools(toolLimits, focused, tools); + return ( +
+
+ + + {replaceKey} +
+ +
+ {tools + .filter((tool) => matches(tool, toolQuery)) + .map((tool) => ( + + onChange(selection, toggleToolLimit(toolLimits, { server: focused, tools }, tool, on)) + } + /> + ))} +
+ +
+ ); + }; + + const nothingAvailable = !isLoading && servers.length === 0 && toolsets.length === 0; + + return ( + !open && setReplacingKey(null)}> + }> + + MCP + {selection.length > 0 && ( + + {selection.length} + + )} + + + + MCP servers + + Pick the servers the model can use, then limit their tools. A server allows every tool until you uncheck + some. + + + {nothingAvailable ? ( +
+ + Add one on the MCP servers page, or send as a key that can reach one + +
+ ) : ( +
+
+
+ +
+ + {isLoading ? ( +
+ +
+ ) : ( +
+ + +
+ )} + +
+
{detail()}
+
+ )} + + }>Done + + + + ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/ModelName.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/ModelName.tsx new file mode 100644 index 00000000000..bbea909711a --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/ModelName.tsx @@ -0,0 +1,74 @@ +"use client"; + +import { Copy, Split } from "lucide-react"; +import { Badge } from "@/components/ui/badge"; +import { Button } from "@/components/ui/button"; +import { HoverCard, HoverCardContent, HoverCardTrigger } from "@/components/ui/hover-card"; +import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; +import { cn } from "@/lib/cva.config"; +import { toast } from "@/lib/toast"; + +export function copyText(value: string, label: string): void { + void navigator.clipboard.writeText(value).then( + () => toast.success(`${label} copied`), + () => toast.error("Unable to copy to clipboard"), + ); +} + +export function CopyIconButton({ value, label }: { value: string; label: string }) { + return ( + + copyText(value, label)} /> + } + > + + + Copy {label} + + ); +} + +export function LongValue({ + value, + detail, + copyLabel, + className, + align = "start", +}: { + value: string; + detail: string; + copyLabel: string; + className?: string; + align?: "start" | "end"; +}) { + return ( + + } + > + {value} + + +
+
+ {detail} + {value} +
+ +
+ + + ); +} + +export function FallbackBadge() { + return ( + + + Fallback + + ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/Pickers.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/Pickers.test.tsx new file mode 100644 index 00000000000..92258ed1428 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/Pickers.test.tsx @@ -0,0 +1,83 @@ +import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { describe, expect, it, vi } from "vitest"; +import { ChecklistMenu, TargetPicker, type OptionGroup } from "./Pickers"; + +const GROUPS: readonly OptionGroup[] = [ + { + label: "openai", + options: [ + { value: "gpt-fast", label: "gpt-fast" }, + { value: "gpt-smart", label: "gpt-smart" }, + ], + }, +]; + +const renderTargetPicker = (onValueChange = vi.fn()) => { + render( + , + ); + return onValueChange; +}; + +describe("TargetPicker", () => { + it("selects the first match when Enter is pressed after searching", async () => { + const user = userEvent.setup(); + const onValueChange = renderTargetPicker(); + + await user.click(screen.getByRole("combobox", { name: "Model" })); + await user.type(await screen.findByPlaceholderText("Search models"), "smart{Enter}"); + + expect(onValueChange).toHaveBeenCalledWith("gpt-smart"); + }); + + it("uses a typed model name that matches nothing when Enter is pressed", async () => { + const user = userEvent.setup(); + const onValueChange = renderTargetPicker(); + + await user.click(screen.getByRole("combobox", { name: "Model" })); + await user.type(await screen.findByPlaceholderText("Search models"), "my-custom-model{Enter}"); + + expect(onValueChange).toHaveBeenCalledWith("my-custom-model"); + }); +}); + +describe("ChecklistMenu", () => { + it("toggles the first match when Enter is pressed after searching", async () => { + const user = userEvent.setup(); + const onValueChange = vi.fn(); + render( + , + ); + + await user.click(screen.getByRole("combobox", { name: "Guardrails" })); + await user.type(await screen.findByPlaceholderText("Search guardrails"), "tox{Enter}"); + + expect(onValueChange).toHaveBeenLastCalledWith(["pii-mask", "toxicity"]); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/Pickers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/Pickers.tsx new file mode 100644 index 00000000000..28ce6fe2f4a --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/Pickers.tsx @@ -0,0 +1,193 @@ +"use client"; + +import { useState, type ReactNode } from "react"; +import { Plus } from "lucide-react"; +import { Badge } from "@/components/ui/badge"; +import { + Combobox, + ComboboxCollection, + ComboboxContent, + ComboboxEmpty, + ComboboxGroup, + ComboboxInput, + ComboboxItem, + ComboboxLabel, + ComboboxList, + ComboboxTrigger, +} from "@/components/ui/combobox"; +import { InputGroupButton } from "@/components/ui/input-group"; +import { Spinner } from "@/components/ui/spinner"; +import { cn } from "@/lib/cva.config"; +import type { PickerOption } from "./usePlaygroundCatalog"; + +export type OptionGroup = Readonly<{ label: string; options: readonly PickerOption[] }>; + +type GroupItems = Readonly<{ value: string; items: readonly string[] }>; + +const toItems = (groups: readonly OptionGroup[]): GroupItems[] => + groups + .filter((group) => group.options.length > 0) + .map((group) => ({ value: group.label, items: group.options.map((option) => option.value) })); + +const isListed = (index: ReadonlyMap, name: string): boolean => + index.has(name) || [...index.values()].some((option) => option.label === name); + +const optionIndex = (groups: readonly OptionGroup[]): ReadonlyMap => + new Map(groups.flatMap((group) => group.options.map((option) => [option.value, option] as const))); + +function GroupedOptions({ index, mono = false }: { index: ReadonlyMap; mono?: boolean }) { + return ( + + {(group: GroupItems) => ( + + {group.value} + + {(value: string) => { + const option = index.get(value); + if (!option) return ; + return ( + + + {option.label} + {option.description && ( + {option.description} + )} + + + ); + }} + + + )} + + ); +} + +function CustomItem({ value }: { value: string }) { + return ( + + + + Use {value} + + + ); +} + +function CountBadge({ count }: { count: number }) { + if (count === 0) return null; + return ( + + {count} + + ); +} + +export function ChecklistMenu({ + label, + icon, + groups, + value, + onValueChange, + isLoading, + emptyText, + defaultOpen = false, +}: { + label: string; + icon: ReactNode; + groups: readonly OptionGroup[]; + value: readonly string[]; + onValueChange: (value: readonly string[]) => void; + isLoading: boolean; + emptyText: string; + defaultOpen?: boolean; +}) { + const index = optionIndex(groups); + const items = toItems(groups); + return ( + onValueChange(next)} + itemToStringLabel={(item: string) => index.get(item)?.label ?? item} + defaultOpen={defaultOpen} + autoHighlight + > + }> + {icon} + {label} + + + + + + {isLoading ? : items.length === 0 && emptyText} + {!isLoading && items.length > 0 && "Nothing matches"} + + + + + ); +} + +export function TargetPicker({ + label, + placeholder, + groups, + value, + onValueChange, + isLoading, + emptyText, + allowCustom = false, + defaultOpen = false, +}: { + label: string; + placeholder: string; + groups: readonly OptionGroup[]; + value: string | null; + onValueChange: (value: string) => void; + isLoading: boolean; + emptyText: string; + allowCustom?: boolean; + defaultOpen?: boolean; +}) { + const [query, setQuery] = useState(""); + const index = optionIndex(groups); + const custom = query.trim(); + const offerCustom = allowCustom && custom !== "" && !isListed(index, custom); + const items = [...toItems(groups), ...(offerCustom ? [{ value: "Custom name", items: [custom] }] : [])]; + const unsetText = isLoading ? `Loading ${label.toLowerCase()}s` : placeholder; + return ( + next && onValueChange(next)} + inputValue={query} + onInputValueChange={setQuery} + onOpenChange={(open) => !open && setQuery("")} + itemToStringLabel={(item: string) => index.get(item)?.label ?? item} + defaultOpen={defaultOpen} + autoHighlight + > + }> + {isLoading && !value ? : null} + + {value ? index.get(value)?.label ?? value : unsetText} + + + + + + {items.length === 0 && !custom ? emptyText : `No matching ${label.toLowerCase()}`} + + + + + ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/PlaygroundChat.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/PlaygroundChat.integration.test.tsx new file mode 100644 index 00000000000..c80f47045c2 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/PlaygroundChat.integration.test.tsx @@ -0,0 +1,658 @@ +import { act, fireEvent, screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { renderWithProviders as render, testQueryClient } from "@/../tests/test-utils"; +import { makeOpenAIChatCompletionRequest } from "@/components/llm_calls/chat_completion"; +import { fetchAvailableModels } from "@/components/llm_calls/fetch_models"; +import { makeOpenAIResponsesRequest } from "@/components/llm_calls/responses_api"; +import { callMCPTool, fetchMCPServers, fetchMCPToolsets, listMCPTools, tagListCall } from "@/components/networking"; +import { setSecureItem } from "@/utils/secureStorage"; +import { fetchClient } from "@/lib/http/api"; +import { PlaygroundChat } from "./PlaygroundChat"; + +vi.mock("@/app/(dashboard)/hooks/useCan", () => ({ default: () => true })); + +vi.mock("@/lib/http/api", async (importOriginal) => { + const actual = await importOriginal(); + return { ...actual, fetchClient: { ...actual.fetchClient, POST: vi.fn() } }; +}); + +vi.mock("@/components/llm_calls/fetch_models", () => ({ fetchAvailableModels: vi.fn() })); + +vi.mock("@/components/llm_calls/chat_completion", () => ({ makeOpenAIChatCompletionRequest: vi.fn() })); + +vi.mock("@/components/llm_calls/responses_api", () => ({ makeOpenAIResponsesRequest: vi.fn() })); + +vi.mock("@/components/networking", () => ({ + tagListCall: vi.fn().mockResolvedValue({}), + vectorStoreListCall: vi.fn().mockResolvedValue({ data: [] }), + getGuardrailsList: vi.fn().mockResolvedValue({ guardrails: [{ guardrail_name: "pii-mask" }] }), + getPoliciesList: vi.fn().mockResolvedValue({ policies: [] }), + fetchMCPServers: vi.fn().mockResolvedValue([]), + fetchMCPToolsets: vi.fn().mockResolvedValue([]), + listMCPTools: vi.fn().mockResolvedValue({ tools: [] }), + callMCPTool: vi.fn(), + getGatewayBaseUrl: () => "http://gateway.test", + getGlobalLitellmHeaderName: () => "Authorization", +})); + +const chatRequest = vi.mocked(makeOpenAIChatCompletionRequest); + +const responsesRequest = vi.mocked(makeOpenAIResponsesRequest); + +const credentialPost = fetchClient.POST as unknown as ReturnType; + +type ChatArgs = Parameters; + +type ResponsesArgs = Parameters; + +const sentMessages = (call: number) => (chatRequest.mock.calls[call] as ChatArgs)[0]; + +const replyWith = + (text: string) => + async (...args: ChatArgs) => { + args[1](text, args[2]); + }; + +const preload = (values: Readonly>) => + Object.entries(values).forEach(([key, value]) => sessionStorage.setItem(key, value)); + +const deferred = () => { + const control: { release: () => void } = { release: () => undefined }; + const promise = new Promise((resolve) => { + control.release = resolve; + }); + return { promise, release: () => control.release() }; +}; + +const mcpTool = (name: string, required: readonly string[] = []) => ({ + name, + inputSchema: { + type: "object", + properties: Object.fromEntries(required.map((field) => [field, { type: "string" }])), + required: [...required], + }, +}); + +const settled = () => waitFor(() => expect(testQueryClient.isFetching()).toBe(0)); + +const renderChat = () => + render(); + +const send = async (user: ReturnType, text: string) => { + await user.type(screen.getByRole("textbox", { name: "Message" }), text); + await user.click(screen.getByRole("button", { name: "Send message" })); +}; + +describe("PlaygroundChat", () => { + beforeEach(() => { + vi.clearAllMocks(); + sessionStorage.clear(); + testQueryClient.clear(); + Element.prototype.scrollIntoView = vi.fn(); + chatRequest.mockReset(); + responsesRequest.mockReset(); + vi.mocked(callMCPTool).mockReset(); + credentialPost.mockReset().mockResolvedValue({ data: {}, response: new Response() }); + vi.mocked(fetchMCPServers).mockReset().mockResolvedValue([]); + vi.mocked(fetchMCPToolsets).mockReset().mockResolvedValue([]); + vi.mocked(listMCPTools).mockReset().mockResolvedValue({ tools: [] }); + vi.mocked(fetchAvailableModels).mockResolvedValue([ + { model_group: "gpt-fast", mode: "chat", providers: ["openai"] }, + ]); + }); + + afterEach(() => { + vi.unstubAllGlobals(); + }); + + it("sends the chosen model and guardrails with the session key", async () => { + const user = userEvent.setup(); + chatRequest.mockImplementation(replyWith("Hi there")); + renderChat(); + + await user.click(screen.getByRole("combobox", { name: "Model" })); + await user.click(await screen.findByRole("option", { name: "gpt-fast" })); + await user.click(screen.getByRole("combobox", { name: "Guardrails" })); + await user.click(await screen.findByRole("option", { name: "pii-mask" })); + await user.keyboard("{Escape}"); + await send(user, "hello"); + + expect(await screen.findByText("Hi there")).toBeInTheDocument(); + const [messages, , model, apiKey, , , , , , , , guardrails] = chatRequest.mock.calls[0] as ChatArgs; + expect(messages).toEqual([{ role: "user", content: "hello" }]); + expect(model).toBe("gpt-fast"); + expect(apiKey).toBe("session-token"); + expect(guardrails).toEqual(["pii-mask"]); + }); + + it("keeps answered turns in the history but never resends a failed one", async () => { + const user = userEvent.setup(); + chatRequest + .mockImplementationOnce(replyWith("Hi there")) + .mockRejectedValueOnce(new Error("rate limited")) + .mockImplementationOnce(replyWith("Back again")); + renderChat(); + + await user.click(screen.getByRole("combobox", { name: "Model" })); + await user.click(await screen.findByRole("option", { name: "gpt-fast" })); + await send(user, "hello"); + expect(await screen.findByText("Hi there")).toBeInTheDocument(); + await send(user, "this one fails"); + expect(await screen.findByText("The request failed")).toBeInTheDocument(); + await send(user, "try again"); + + await waitFor(() => expect(chatRequest).toHaveBeenCalledTimes(3)); + expect(sentMessages(2)).toEqual([ + { role: "user", content: "hello" }, + { role: "assistant", content: "Hi there" }, + { role: "user", content: "try again" }, + ]); + }); + + it("sends every saved request setting with the chat request", async () => { + const user = userEvent.setup(); + chatRequest.mockImplementation(replyWith("Hi there")); + vi.mocked(fetchMCPServers).mockResolvedValue([{ server_id: "srv-1", alias: "docs" }]); + preload({ + selectedModel: "gpt-fast", + selectedTags: JSON.stringify(["team-a"]), + selectedVectorStores: JSON.stringify(["vs-1"]), + selectedPolicies: JSON.stringify(["strict"]), + selectedMCPServers: JSON.stringify(["srv-1"]), + streamingEnabled: "false", + }); + setSecureItem("customHeaders", JSON.stringify([["x-team", "blue"]])); + renderChat(); + + await waitFor(() => expect(fetchMCPServers).toHaveBeenCalled()); + await send(user, "hello"); + + await waitFor(() => expect(chatRequest).toHaveBeenCalledTimes(1)); + const args = chatRequest.mock.calls[0] as ChatArgs; + expect(args[2]).toBe("gpt-fast"); + expect(args[4]).toEqual(["team-a"]); + expect(args[10]).toEqual(["vs-1"]); + expect(args[12]).toEqual(["strict"]); + expect(args[13]).toEqual(["srv-1"]); + expect(args[25]).toBe(false); + expect(args[26]).toMatchObject({ "x-team": "blue" }); + }); + + it("sends a custom model and tag that are prefixes of listed ones", async () => { + const user = userEvent.setup(); + chatRequest.mockImplementation(replyWith("Hi there")); + vi.mocked(tagListCall).mockResolvedValueOnce({ + production: { name: "production", models: [], created_at: "", updated_at: "" }, + }); + renderChat(); + + await user.click(screen.getByRole("combobox", { name: "Model" })); + await user.type(await screen.findByPlaceholderText("Search models"), "gpt-f"); + await user.click(await screen.findByRole("option", { name: "Use gpt-f" })); + await user.click(screen.getByRole("button", { name: "Settings" })); + await user.type(await screen.findByLabelText("Tags"), "prod"); + await user.click(await screen.findByRole("option", { name: "Add prod" })); + await send(user, "hello"); + + await waitFor(() => expect(chatRequest).toHaveBeenCalledTimes(1)); + const args = chatRequest.mock.calls[0] as ChatArgs; + expect(args[2]).toBe("gpt-f"); + expect(args[4]).toEqual(["prod"]); + }); + + it("sends a temperature typed with two decimals", async () => { + const user = userEvent.setup(); + chatRequest.mockImplementation(replyWith("Hi there")); + preload({ selectedModel: "gpt-fast" }); + renderChat(); + + await user.click(screen.getByRole("button", { name: "Settings" })); + await user.click(await screen.findByRole("button", { name: "Advanced" })); + await user.type(screen.getByLabelText("Temperature value"), "0.05"); + await user.keyboard("{Escape}"); + await send(user, "hello"); + + await waitFor(() => expect(chatRequest).toHaveBeenCalledTimes(1)); + expect((chatRequest.mock.calls[0] as ChatArgs)[16]).toBe(0.05); + }); + + it("keeps the saved MCP selection when the server list cannot be fetched", async () => { + const user = userEvent.setup(); + chatRequest.mockImplementation(replyWith("Hi there")); + vi.mocked(fetchMCPServers).mockRejectedValue(new Error("forbidden")); + preload({ selectedModel: "gpt-fast", selectedMCPServers: JSON.stringify(["srv-1"]) }); + renderChat(); + + await waitFor(() => expect(fetchMCPServers).toHaveBeenCalled()); + await settled(); + await send(user, "hello"); + + await waitFor(() => expect(chatRequest).toHaveBeenCalledTimes(1)); + expect((chatRequest.mock.calls[0] as ChatArgs)[13]).toEqual(["srv-1"]); + }); + + it("drops a saved model that the current key cannot use", async () => { + const user = userEvent.setup(); + chatRequest.mockImplementation(replyWith("Hi there")); + preload({ selectedModel: "retired-model" }); + renderChat(); + + await waitFor(() => expect(fetchAvailableModels).toHaveBeenCalled()); + await waitFor(() => expect(screen.getByRole("combobox", { name: "Model" })).toHaveTextContent("Select a model")); + await user.click(screen.getByRole("combobox", { name: "Model" })); + await user.click(await screen.findByRole("option", { name: "gpt-fast" })); + await send(user, "hello"); + + await waitFor(() => expect(chatRequest).toHaveBeenCalledTimes(1)); + expect((chatRequest.mock.calls[0] as ChatArgs)[2]).toBe("gpt-fast"); + }); + + it("calls a toolset tool on the server that backs it", async () => { + const user = userEvent.setup(); + vi.mocked(fetchMCPToolsets).mockResolvedValue([ + { toolset_id: "ts-1", toolset_name: "Research", tools: [{ server_id: "srv-2", tool_name: "search" }] }, + ]); + vi.mocked(listMCPTools).mockResolvedValue({ + tools: [{ name: "search", inputSchema: { type: "object", properties: {} } }], + }); + vi.mocked(callMCPTool).mockResolvedValue({ content: [{ type: "text", text: "found it" }] }); + preload({ endpointType: "mcp", selectedMCPServers: JSON.stringify(["toolset:ts-1"]) }); + renderChat(); + + await user.click(await screen.findByRole("combobox", { name: "Tool" })); + await user.click(await screen.findByRole("option", { name: /search/ })); + await user.click(screen.getByRole("button", { name: "Send message" })); + + await waitFor(() => expect(callMCPTool).toHaveBeenCalledTimes(1)); + expect(listMCPTools).toHaveBeenCalledWith("session-token", "srv-2"); + expect(listMCPTools).not.toHaveBeenCalledWith("session-token", "toolset:ts-1"); + expect(vi.mocked(callMCPTool).mock.calls[0]?.slice(1, 3)).toEqual(["srv-2", "search"]); + }); + + it("offers toolsets as targets and lists only the toolset's own tools", async () => { + const user = userEvent.setup(); + vi.mocked(fetchMCPToolsets).mockResolvedValue([ + { + toolset_id: "ts-1", + toolset_name: "Research", + tools: [ + { server_id: "srv-a", tool_name: "search" }, + { server_id: "srv-b", tool_name: "fetch" }, + { server_id: "srv-b", tool_name: "search" }, + ], + }, + ]); + vi.mocked(listMCPTools).mockImplementation(async (_key: string, server: string) => ({ + tools: + server === "srv-a" ? [mcpTool("search"), mcpTool("fetch", ["a_arg"])] : [mcpTool("search"), mcpTool("fetch")], + })); + vi.mocked(callMCPTool).mockResolvedValue({ content: [{ type: "text", text: "fetched" }] }); + preload({ endpointType: "mcp" }); + renderChat(); + + await user.click(await screen.findByRole("combobox", { name: "Server" })); + await user.click(await screen.findByRole("option", { name: /Research/ })); + await settled(); + await user.click(screen.getByRole("combobox", { name: "Tool" })); + await waitFor(() => + expect(screen.getAllByRole("option").map((option) => option.textContent)).toEqual(["search", "fetch"]), + ); + await user.click(screen.getByRole("option", { name: "fetch" })); + await user.click(screen.getByRole("button", { name: "Send message" })); + + await waitFor(() => expect(callMCPTool).toHaveBeenCalledTimes(1)); + expect(vi.mocked(callMCPTool).mock.calls[0]?.slice(1, 4)).toEqual(["srv-b", "fetch", {}]); + }); + + it("says the tools could not be loaded when the server's tool list fails", async () => { + const user = userEvent.setup(); + vi.mocked(fetchMCPServers).mockResolvedValue([{ server_id: "srv-1", alias: "docs" }]); + vi.mocked(listMCPTools).mockResolvedValue({ error: "server unreachable" }); + preload({ endpointType: "mcp", selectedMCPServers: JSON.stringify(["srv-1"]) }); + renderChat(); + + await settled(); + await user.click(screen.getByRole("combobox", { name: "Tool" })); + + expect(await screen.findByText("Unable to load tools for this server")).toBeInTheDocument(); + }); + + it("keeps the saved MCP selection when the toolset list cannot be fetched", async () => { + const user = userEvent.setup(); + chatRequest.mockImplementation(replyWith("Hi there")); + vi.mocked(fetchMCPServers).mockResolvedValue([{ server_id: "srv-1", alias: "docs" }]); + vi.mocked(fetchMCPToolsets).mockRejectedValue(new Error("database not connected")); + preload({ selectedModel: "gpt-fast", selectedMCPServers: JSON.stringify(["srv-1", "srv-gone"]) }); + renderChat(); + + await waitFor(() => expect(fetchMCPToolsets).toHaveBeenCalled()); + await settled(); + await send(user, "hello"); + + await waitFor(() => expect(chatRequest).toHaveBeenCalledTimes(1)); + expect((chatRequest.mock.calls[0] as ChatArgs)[13]).toEqual(["srv-1", "srv-gone"]); + }); + + it("drops a saved model that is missing from an already loaded model list", async () => { + const first = renderChat(); + await waitFor(() => expect(fetchAvailableModels).toHaveBeenCalled()); + await settled(); + first.unmount(); + sessionStorage.setItem("selectedModel", "retired-model"); + + renderChat(); + + await waitFor(() => expect(screen.getByRole("combobox", { name: "Model" })).toHaveTextContent("Select a model")); + expect(fetchAvailableModels).toHaveBeenCalledTimes(1); + }); + + it("drops a request that is still running when a new chat starts", async () => { + const user = userEvent.setup(); + const first = deferred(); + chatRequest.mockImplementationOnce(async (...args: ChatArgs) => { + await first.promise; + args[1]("late reply", args[2]); + }); + preload({ selectedModel: "gpt-fast" }); + renderChat(); + + await send(user, "first"); + expect(await screen.findByRole("button", { name: "Stop" })).toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "New chat" })); + await act(async () => { + first.release(); + await chatRequest.mock.results[0]?.value; + }); + + expect(chatRequest).toHaveBeenCalledTimes(1); + expect(screen.queryByRole("button", { name: "Stop" })).not.toBeInTheDocument(); + expect(screen.queryByText("late reply")).not.toBeInTheDocument(); + }); + + it("keeps a stopped request's response id out of the next request", async () => { + const user = userEvent.setup(); + const first = deferred(); + responsesRequest + .mockImplementationOnce(async (...args: ResponsesArgs) => { + await first.promise; + args[15]?.("resp-stopped"); + }) + .mockImplementationOnce(async (...args: ResponsesArgs) => { + args[15]?.("resp-fresh"); + args[1]("assistant", "fresh reply", args[2]); + }) + .mockImplementationOnce(async (...args: ResponsesArgs) => { + args[1]("assistant", "third reply", args[2]); + }); + preload({ endpointType: "responses", selectedModel: "gpt-fast", useApiSessionManagement: "true" }); + renderChat(); + + await send(user, "first"); + await user.click(await screen.findByRole("button", { name: "Stop" })); + await send(user, "second"); + expect(await screen.findByText("fresh reply")).toBeInTheDocument(); + await act(async () => { + first.release(); + await responsesRequest.mock.results[0]?.value; + }); + await send(user, "third"); + + await waitFor(() => expect(responsesRequest).toHaveBeenCalledTimes(3)); + expect((responsesRequest.mock.calls[2] as ResponsesArgs)[14]).toBe("resp-fresh"); + }); + + it("ignores a stopped request that finishes after a newer one started", async () => { + const user = userEvent.setup(); + const first = deferred(); + const second = deferred(); + chatRequest + .mockImplementationOnce(async (...args: ChatArgs) => { + await first.promise; + args[1]("late reply", args[2]); + }) + .mockImplementationOnce(async (...args: ChatArgs) => { + await second.promise; + args[1]("fresh reply", args[2]); + }); + preload({ selectedModel: "gpt-fast" }); + renderChat(); + + await send(user, "first"); + await user.click(await screen.findByRole("button", { name: "Stop" })); + await send(user, "second"); + first.release(); + + await waitFor(() => expect(chatRequest).toHaveBeenCalledTimes(2)); + expect(screen.getByRole("button", { name: "Stop" })).toBeInTheDocument(); + expect(screen.queryByText("late reply")).not.toBeInTheDocument(); + second.release(); + expect(await screen.findByText("fresh reply")).toBeInTheDocument(); + expect(screen.queryByText("late reply")).not.toBeInTheDocument(); + }); + + it("only says no guardrails ran when the gateway reported its headers", async () => { + const user = userEvent.setup(); + vi.stubGlobal("matchMedia", (query: string) => ({ + matches: true, + media: query, + addEventListener: () => undefined, + removeEventListener: () => undefined, + })); + chatRequest + .mockImplementationOnce(replyWith("unreported reply")) + .mockImplementationOnce(async (...args: ChatArgs) => { + args[27]?.(new Headers({ "x-litellm-call-id": "call-reported" })); + args[1]("reported reply", args[2]); + }); + preload({ selectedModel: "gpt-fast", selectedGuardrails: JSON.stringify(["pii-mask"]) }); + renderChat(); + + await send(user, "first"); + expect(await screen.findByText("unreported reply")).toBeInTheDocument(); + expect(screen.getByText("This endpoint does not report which guardrails or policies ran")).toBeInTheDocument(); + expect(screen.queryByText("No guardrails or policies ran on this request")).not.toBeInTheDocument(); + + await send(user, "second"); + expect(await screen.findByText("call-reported")).toBeInTheDocument(); + expect(screen.getByText("No guardrails or policies ran on this request")).toBeInTheDocument(); + }); + + it("does not claim policies were skipped when the gateway's own headers are hidden", async () => { + const user = userEvent.setup(); + vi.stubGlobal("matchMedia", (query: string) => ({ + matches: true, + media: query, + addEventListener: () => undefined, + removeEventListener: () => undefined, + })); + chatRequest.mockImplementationOnce(async (...args: ChatArgs) => { + args[27]?.(new Headers({ "content-type": "application/json" })); + args[1]("cross-origin reply", args[2]); + }); + preload({ selectedModel: "gpt-fast" }); + renderChat(); + + await send(user, "hello"); + expect(await screen.findByText("cross-origin reply")).toBeInTheDocument(); + expect( + screen.getByText("Policies are only reported when the dashboard is opened from the proxy"), + ).toBeInTheDocument(); + expect(screen.queryByText(/policies ran on this request/)).not.toBeInTheDocument(); + }); + + it("keeps a stopped request's gateway details off the newer reply", async () => { + const user = userEvent.setup(); + vi.stubGlobal("matchMedia", (query: string) => ({ + matches: true, + media: query, + addEventListener: () => undefined, + removeEventListener: () => undefined, + })); + const first = deferred(); + chatRequest + .mockImplementationOnce(async (...args: ChatArgs) => { + args[27]?.(new Headers({ "x-litellm-call-id": "call-stopped" })); + await first.promise; + }) + .mockImplementationOnce(async (...args: ChatArgs) => { + args[27]?.(new Headers({ "x-litellm-call-id": "call-fresh" })); + args[1]("fresh reply", args[2]); + }); + preload({ selectedModel: "gpt-fast" }); + renderChat(); + + await send(user, "first"); + await user.click(await screen.findByRole("button", { name: "Stop" })); + await send(user, "second"); + expect(await screen.findByText("call-fresh")).toBeInTheDocument(); + await act(async () => { + first.release(); + await chatRequest.mock.results[0]?.value; + }); + + expect(screen.queryByRole("button", { name: "Stop" })).not.toBeInTheDocument(); + expect(screen.getByText("call-fresh")).toBeInTheDocument(); + expect(screen.queryByText("call-stopped")).not.toBeInTheDocument(); + }); + + it("sends once when Enter is pressed twice before the first request starts", async () => { + const user = userEvent.setup(); + const first = deferred(); + chatRequest.mockImplementation(async (...args: ChatArgs) => { + await first.promise; + args[1]("Hi there", args[2]); + }); + preload({ selectedModel: "gpt-fast" }); + renderChat(); + + const message = screen.getByRole("textbox", { name: "Message" }); + await user.type(message, "hello"); + fireEvent.keyDown(message, { key: "Enter" }); + fireEvent.keyDown(message, { key: "Enter" }); + + await waitFor(() => expect(chatRequest).toHaveBeenCalled()); + await act(async () => { + first.release(); + await chatRequest.mock.results[0]?.value; + }); + expect(chatRequest).toHaveBeenCalledTimes(1); + expect(screen.getAllByText("hello")).toHaveLength(1); + }); + + it("drops a chat attachment when the endpoint cannot use it", async () => { + const user = userEvent.setup(); + preload({ selectedModel: "gpt-fast" }); + renderChat(); + + fireEvent.drop(screen.getByRole("textbox", { name: "Message" }), { + dataTransfer: { files: [new File(["%PDF"], "notes.pdf", { type: "application/pdf" })] }, + }); + expect(await screen.findByRole("button", { name: "Remove notes.pdf" })).toBeInTheDocument(); + + await user.click(screen.getByRole("button", { name: "Endpoint" })); + await user.click(await screen.findByRole("menuitemradio", { name: "/v1/audio/transcriptions" })); + + await waitFor(() => + expect(screen.getByRole("button", { name: "Endpoint" })).toHaveTextContent("audio/transcriptions"), + ); + expect(screen.queryByRole("button", { name: "Remove notes.pdf" })).not.toBeInTheDocument(); + }); + + it("keeps a stopped request's code output off the newer reply", async () => { + const user = userEvent.setup(); + const first = deferred(); + responsesRequest + .mockImplementationOnce(async (...args: ResponsesArgs) => { + await first.promise; + args[18]?.({ code: "print('stopped run')", containerId: "c-1", annotations: [] }); + }) + .mockImplementationOnce(async (...args: ResponsesArgs) => { + args[1]("assistant", "fresh reply", args[2]); + }); + preload({ endpointType: "responses", selectedModel: "gpt-fast" }); + renderChat(); + + await send(user, "first"); + await user.click(await screen.findByRole("button", { name: "Stop" })); + await send(user, "second"); + expect(await screen.findByText("fresh reply")).toBeInTheDocument(); + await act(async () => { + first.release(); + await responsesRequest.mock.results[0]?.value; + }); + + expect(screen.queryByText("Python Code Executed")).not.toBeInTheDocument(); + }); + + it("lets a user connect a key for a BYOK server picked for a direct tool call", async () => { + const user = userEvent.setup(); + vi.mocked(fetchMCPServers).mockResolvedValue([ + { server_id: "srv-byok", alias: "Linear", is_byok: true, has_user_credential: false }, + ]); + preload({ endpointType: "mcp", selectedMCPServers: JSON.stringify(["srv-byok"]) }); + renderChat(); + + await user.click(await screen.findByRole("button", { name: "Connect key" })); + await user.type(await screen.findByLabelText("Linear API key"), "lin-key-1"); + await user.click(screen.getByRole("button", { name: "Connect" })); + + await waitFor(() => expect(credentialPost).toHaveBeenCalledTimes(1)); + expect(credentialPost.mock.calls[0]?.[1]).toMatchObject({ + params: { path: { server_id: "srv-byok" } }, + body: { credential: "lin-key-1" }, + }); + await waitFor(() => expect(fetchMCPServers).toHaveBeenCalledTimes(2)); + }); + + it("lets a user replace the key of a connected BYOK server whose tools fail to load", async () => { + const user = userEvent.setup(); + vi.mocked(fetchMCPServers).mockResolvedValue([ + { server_id: "srv-byok", alias: "Linear", is_byok: true, has_user_credential: true }, + ]); + vi.mocked(listMCPTools).mockResolvedValue({ error: "401 from upstream" }); + preload({ selectedModel: "gpt-fast", selectedMCPServers: JSON.stringify(["srv-byok"]) }); + renderChat(); + + await settled(); + await user.click(screen.getByRole("button", { name: "MCP" })); + await user.click(await screen.findByRole("button", { name: "Replace key" })); + await user.type(await screen.findByLabelText("Linear API key"), "lin-key-2"); + await user.click(screen.getByRole("button", { name: "Connect" })); + + await waitFor(() => expect(credentialPost).toHaveBeenCalledTimes(1)); + expect(credentialPost.mock.calls[0]?.[1]).toMatchObject({ body: { credential: "lin-key-2" } }); + await waitFor(() => + expect(vi.mocked(listMCPTools).mock.calls.filter(([, server]) => server === "srv-byok").length).toBeGreaterThan( + 1, + ), + ); + }); + + it("leaves replace-key mode on cancel or when the MCP dialog closes", async () => { + const user = userEvent.setup(); + vi.mocked(fetchMCPServers).mockResolvedValue([ + { server_id: "srv-byok", alias: "Linear", is_byok: true, has_user_credential: true }, + ]); + vi.mocked(listMCPTools).mockResolvedValue({ error: "401 from upstream" }); + preload({ selectedModel: "gpt-fast", selectedMCPServers: JSON.stringify(["srv-byok"]) }); + renderChat(); + + await settled(); + await user.click(screen.getByRole("button", { name: "MCP" })); + await user.click(await screen.findByRole("button", { name: "Replace key" })); + expect(await screen.findByLabelText("Linear API key")).toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "Cancel" })); + expect(screen.queryByLabelText("Linear API key")).not.toBeInTheDocument(); + + await user.click(screen.getByRole("button", { name: "Replace key" })); + expect(await screen.findByLabelText("Linear API key")).toBeInTheDocument(); + await user.keyboard("{Escape}"); + await waitFor(() => expect(screen.queryByLabelText("Linear API key")).not.toBeInTheDocument()); + await user.click(screen.getByRole("button", { name: "MCP" })); + + expect(await screen.findByRole("button", { name: "Replace key" })).toBeInTheDocument(); + expect(screen.queryByLabelText("Linear API key")).not.toBeInTheDocument(); + expect(credentialPost).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/PlaygroundChat.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/PlaygroundChat.tsx new file mode 100644 index 00000000000..5b2de80315d --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/PlaygroundChat.tsx @@ -0,0 +1,489 @@ +"use client"; + +import { useRef, useState } from "react"; +import { ChevronDown, PanelRight, RotateCcw, SquarePen } from "lucide-react"; +import useCan from "@/app/(dashboard)/hooks/useCan"; +import { generateCodeSnippet } from "@/components/chat_ui/CodeSnippets"; +import { EndpointType, getEndpointType } from "@/components/chat_ui/mode_endpoint_mapping"; +import type { MessageType } from "@/components/chat_ui/types"; +import type { ModelGroup } from "@/components/llm_calls/fetch_models"; +import { customHeadersFromPairs } from "@/components/llm_calls/request_headers"; +import type { MCPToolArgumentsFormRef } from "@/components/mcp_tools/MCPToolArgumentsForm"; +import type { MCPServer, MCPTool, MCPToolset } from "@/components/mcp_tools/types"; +import { Button } from "@/components/ui/button"; +import { ButtonGroup } from "@/components/ui/button-group"; +import { + DropdownMenu, + DropdownMenuContent, + DropdownMenuItem, + DropdownMenuTrigger, +} from "@/components/ui/dropdown-menu"; +import { Sheet, SheetContent, SheetTitle } from "@/components/ui/sheet"; +import { toast } from "@/lib/toast"; +import { useChatHistory } from "../../hooks/useChatHistory"; +import { useCodeInterpreter } from "../../hooks/useCodeInterpreter"; +import { filterModelsForEndpoint, isModelCompatibleWithEndpoint } from "../chat_ui/EndpointUtils"; +import RealtimePlayground from "../chat_ui/RealtimePlayground"; +import { Composer } from "./Composer"; +import { EndpointMenu, VoiceMenu } from "./ComposerControls"; +import { Conversation, formatCost } from "./Conversation"; +import { capabilitiesFor, CODE_INTERPRETER_SUGGESTIONS, SUGGESTIONS } from "./endpointCapabilities"; +import { InspectorPanel, type CodeSdk } from "./Inspector"; +import { ByokKeyButton } from "./ByokConnect"; +import { McpDialog, serverName } from "./McpDialog"; +import { TargetPicker, type OptionGroup } from "./Pickers"; +import { TOOLSET_PREFIX, withEndpoint, withReachableMcp } from "./playgroundSettings"; +import { SettingsPopover } from "./SettingsPopover"; +import { + useGuardrailOptions, + useMcpCatalog, + usePlaygroundAgents, + usePlaygroundModels, + usePolicyOptions, + useServerTools, + useTagOptions, + useVectorStoreOptions, + type PickerOption, + type ServerTools, +} from "./usePlaygroundCatalog"; +import { usePlaygroundSend, type SendDeps } from "./usePlaygroundSend"; +import { usePlaygroundSettings } from "./usePlaygroundSettings"; +import { useIsWide } from "./useIsWide"; + +export type ProxySettings = Readonly<{ PROXY_BASE_URL?: string; LITELLM_UI_API_DOC_BASE_URL?: string | null }>; + +const POLICY_PREFIX = "policy:"; + +const providerOf = (model: ModelGroup): string => model.providers?.[0] ?? "other"; + +const modelGroups = (models: readonly ModelGroup[]): readonly OptionGroup[] => + [...new Set(models.map(providerOf))].sort().map((provider) => ({ + label: provider, + options: models + .filter((model) => providerOf(model) === provider) + .map((model) => ({ value: model.model_group, label: model.model_group })), + })); + +const modelEmptyText = (failed: boolean, hasKey: boolean): string => { + if (failed) return "Unable to load models for this key"; + if (!hasKey) return "Enter a virtual key to load models"; + return "No models available for this endpoint"; +}; + +const reachableMcp = (servers: readonly MCPServer[], toolsets: readonly MCPToolset[]): ReadonlySet => + new Set([ + ...servers.map((server) => server.server_id), + ...toolsets.map((toolset) => `${TOOLSET_PREFIX}${toolset.toolset_id}`), + ]); + +const toolsetFor = (selection: string | null, toolsets: readonly MCPToolset[]): MCPToolset | null => + toolsets.find((toolset) => `${TOOLSET_PREFIX}${toolset.toolset_id}` === selection) ?? null; + +const backingServers = (selection: string | null, toolset: MCPToolset | null): readonly string[] => { + if (toolset) return [...new Set(toolset.tools.map((tool) => tool.server_id))]; + return selection && !selection.startsWith(TOOLSET_PREFIX) ? [selection] : []; +}; + +type DirectTarget = Readonly<{ server: string; tool: MCPTool }>; + +const toolsetTargets = (toolset: MCPToolset, loaded: Readonly>): DirectTarget[] => + toolset.tools + .flatMap((entry) => { + const tool = loaded[entry.server_id]?.tools.find((candidate) => candidate.name === entry.tool_name); + return tool ? [{ server: entry.server_id, tool }] : []; + }) + .filter((target, index, targets) => targets.findIndex((other) => other.tool.name === target.tool.name) === index); + +const directTargets = ( + servers: readonly string[], + toolset: MCPToolset | null, + loaded: Readonly>, +): readonly DirectTarget[] => { + if (toolset) return toolsetTargets(toolset, loaded); + return servers.flatMap((server) => (loaded[server]?.tools ?? []).map((tool) => ({ server, tool }))); +}; + +const prefixed = (options: readonly PickerOption[]): readonly PickerOption[] => + options.map((option) => ({ ...option, value: `${POLICY_PREFIX}${option.value}` })); + +const lastTracedReply = (messages: readonly MessageType[]): number => + messages.findLastIndex((message) => message.role === "assistant" && message.trace !== undefined); + +const totalCost = (messages: readonly MessageType[]): number => + messages.reduce((sum, message) => sum + (message.usage?.cost ?? message.trace?.cost ?? 0), 0); + +const conversationSummary = (messages: readonly MessageType[]): string => { + const requests = messages.filter((message) => message.role === "user").length; + if (requests === 0) return "New chat"; + const label = requests === 1 ? "1 request" : `${requests} requests`; + const cost = totalCost(messages); + return cost > 0 ? `${label}, ${formatCost(cost)}` : label; +}; + +export function PlaygroundChat({ + accessToken, + disabledPersonalKeyCreation, + proxySettings, +}: { + accessToken: string | null; + disabledPersonalKeyCreation: boolean; + proxySettings: ProxySettings | undefined; +}) { + const { settings: stored, update, reset } = usePlaygroundSettings(disabledPersonalKeyCreation); + const canViewPolicies = useCan("viewPolicies"); + const apiKey = (stored.keySource === "session" ? accessToken : stored.virtualKey.trim()) || null; + const mcp = useMcpCatalog(apiKey); + const settings = withReachableMcp(stored, mcp.isReady ? reachableMcp(mcp.servers, mcp.toolsets) : null); + const capabilities = capabilitiesFor(settings.endpoint); + const chat = useChatHistory({ simplified: false }); + const codeInterpreter = useCodeInterpreter(); + const toolArgsForm = useRef(null); + const [inspected, setInspected] = useState(null); + const [sheetOpen, setSheetOpen] = useState(false); + const isWide = useIsWide(); + + const models = usePlaygroundModels(apiKey); + const [checkedModels, setCheckedModels] = useState(undefined); + if (models.data && models.data !== checkedModels) { + setCheckedModels(models.data); + const loaded = models.data; + update((current) => + current.model && !loaded.some((option) => option.model_group === current.model) + ? { ...current, model: null } + : current, + ); + } + const agents = usePlaygroundAgents(apiKey, settings.proxyBaseUrl, settings.endpoint === EndpointType.A2A_AGENTS); + const guardrails = useGuardrailOptions(accessToken); + const policies = usePolicyOptions(accessToken, canViewPolicies); + const vectorStores = useVectorStoreOptions(accessToken); + const tags = useTagOptions(accessToken); + + const directSelection = settings.endpoint === EndpointType.MCP ? settings.mcpSelection[0] ?? null : null; + const directToolset = toolsetFor(directSelection, mcp.toolsets); + const directServers = backingServers(directSelection, directToolset); + const directServerTools = useServerTools(apiKey, directServers); + const directByok = directToolset ? undefined : mcp.servers.find((server) => server.server_id === directSelection); + const targets = directTargets(directServers, directToolset, directServerTools); + const directTarget = targets.find((target) => target.tool.name === settings.directTool) ?? null; + const directTool = directTarget?.tool ?? null; + const toolsFailed = directServers.some((server) => directServerTools[server]?.failed); + const customHeaders = customHeadersFromPairs(settings.headerPairs); + + const sendDeps: SendDeps = { + settings, + capabilities, + apiKey, + signedIn: Boolean(accessToken), + customHeaders, + mcpServers: mcp.servers, + mcpToolsets: mcp.toolsets, + directCall: () => (directTarget ? { serverId: directTarget.server, tool: directTarget.tool.name } : null), + toolArgsForm, + chat, + codeInterpreter, + }; + const draft = usePlaygroundSend(sendDeps); + + const messages = chat.chatHistory; + const inspectedIndex = inspected ?? lastTracedReply(messages); + const inspectedMessage = inspectedIndex >= 0 ? messages[inspectedIndex] ?? null : null; + const modelList = models.data ?? []; + + const selectModel = (model: string) => { + const info = modelList.find((option) => option.model_group === model); + const switchTo = + info?.mode && !isModelCompatibleWithEndpoint(info, settings.endpoint) ? getEndpointType(info.mode) : null; + update((current) => ({ ...(switchTo ? withEndpoint(current, switchTo) : current), model })); + }; + + const inspect = (index: number) => { + setInspected(index); + if (!isWide) setSheetOpen(true); + }; + + const clearChat = () => { + chat.clearChatHistory(); + draft.clearDraft(); + codeInterpreter.clearResult(); + setInspected(null); + }; + + const newChat = () => { + clearChat(); + toast.success("Chat history cleared."); + }; + + const resetAll = () => { + clearChat(); + reset(); + toast.success("Chat and settings reset"); + }; + + const code = (sdk: CodeSdk): string => { + const snippet: Parameters[0] = { + apiKeySource: "custom", + accessToken: null, + apiKey: "", + inputMessage: draft.text, + chatHistory: messages, + selectedTags: [...settings.tags], + selectedVectorStores: [...settings.vectorStores], + selectedGuardrails: [...settings.guardrails], + selectedPolicies: [...settings.policies], + selectedMCPServers: [...settings.mcpSelection], + mcpServers: [...mcp.servers], + mcpServerToolRestrictions: Object.fromEntries( + Object.entries(settings.toolLimits).map(([server, tools]) => [server, [...tools]]), + ), + endpointType: settings.endpoint, + selectedModel: settings.model ?? undefined, + selectedSdk: sdk, + selectedVoice: settings.voice, + proxySettings, + customHeaders, + }; + return generateCodeSnippet(snippet); + }; + + const targetRow = ( + <> + {capabilities.target === "model" && ( + + )} + {capabilities.target === "agent" && ( + ({ + value: agent.agent_name, + label: agent.agent_name || agent.agent_id, + description: agent.agent_card_params?.description, + })), + }, + ]} + value={settings.agent} + onValueChange={(agent) => update({ agent })} + isLoading={agents.isLoading} + emptyText="No agents available for this key" + /> + )} + {capabilities.target === "mcpTool" && ( + <> + ({ value: server.server_id, label: serverName(server) })), + }, + { + label: "Toolsets", + options: mcp.toolsets.map((toolset) => ({ + value: `${TOOLSET_PREFIX}${toolset.toolset_id}`, + label: toolset.toolset_name, + description: toolset.description, + })), + }, + ]} + value={directSelection} + onValueChange={(server) => update({ mcpSelection: [server], directTool: null })} + isLoading={mcp.isLoading} + emptyText="No MCP servers available for this key" + /> + {directByok?.is_byok && ( + + )} + {directSelection && ( + ({ + value: tool.name, + label: tool.name, + description: tool.description, + })), + }, + ]} + value={settings.directTool} + onValueChange={(tool) => update({ directTool: tool })} + isLoading={directServers.some((server) => directServerTools[server]?.isLoading)} + emptyText={toolsFailed ? "Unable to load tools for this server" : "This server has no tools"} + /> + )} + + )} + update((current) => withEndpoint(current, endpoint))} + /> + {capabilities.voice && update({ voice })} />} + + ); + + const guardrailGroups: readonly OptionGroup[] = [ + { label: "Guardrails", options: guardrails.data ?? [] }, + ...(canViewPolicies ? [{ label: "Policies", options: prefixed(policies.data ?? []) }] : []), + ]; + + const composer = ( + update({ mcpSelection, toolLimits })} + onCredentialSaved={mcp.refreshCredentials} + /> + } + vectorStores={{ + groups: [{ label: "Vector stores", options: vectorStores.data ?? [] }], + value: settings.vectorStores, + onValueChange: (value) => update({ vectorStores: value }), + isLoading: vectorStores.isLoading, + }} + guardrails={{ + groups: guardrailGroups, + value: [...settings.guardrails, ...settings.policies.map((policy) => `${POLICY_PREFIX}${policy}`)], + onValueChange: (value) => + update({ + guardrails: value.filter((item) => !item.startsWith(POLICY_PREFIX)), + policies: value + .filter((item) => item.startsWith(POLICY_PREFIX)) + .map((item) => item.slice(POLICY_PREFIX.length)), + }), + isLoading: guardrails.isLoading || policies.isLoading, + }} + codeInterpreter={{ enabled: codeInterpreter.enabled, onToggle: codeInterpreter.setEnabled }} + settingsControl={ + + } + sendDisabled={draft.isLoading} + /> + ); + + const suggestions = + settings.endpoint === EndpointType.RESPONSES && codeInterpreter.enabled + ? CODE_INTERPRETER_SUGGESTIONS + : SUGGESTIONS[settings.endpoint] ?? []; + + const inspector = ; + + return ( +
+
+
+ {conversationSummary(messages)} + + {!isWide && ( + + )} + + + + } + > + + + + + + New chat, keep settings + + + + Reset chat and settings + + + + + +
+ {settings.endpoint === EndpointType.REALTIME ? ( +
+ 0 ? [...settings.guardrails] : undefined} + /> +
+ ) : ( + + )} +
{composer}
+
+ {isWide ? ( + + ) : ( + + + Inspector + {inspector} + + + )} +
+ ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/ReplyContent.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/ReplyContent.tsx new file mode 100644 index 00000000000..5733ede7168 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/ReplyContent.tsx @@ -0,0 +1,92 @@ +"use client"; + +import ReactMarkdown from "react-markdown"; +import { Prism as SyntaxHighlighter } from "react-syntax-highlighter"; +import { coy } from "react-syntax-highlighter/dist/esm/styles/prism"; +import remarkGfm from "remark-gfm"; +import type { MessageType } from "@/components/chat_ui/types"; +import { ScrollArea, ScrollBar } from "@/components/ui/scroll-area"; +import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; +import { useSyntaxTheme } from "@/hooks/useSyntaxTheme"; +import AudioRenderer from "../chat_ui/AudioRenderer"; + +const textOf = (message: MessageType): string => (typeof message.content === "string" ? message.content : ""); + +function CodeBlock({ language, code }: { language: string; code: string }) { + const theme = useSyntaxTheme(coy); + return ( + + {code} + + ); +} + +export function Markdown({ children }: { children: string }) { + return ( +

, + ul: ({ node: _node, ...props }) =>