From 492251c7bc7bd76baf25512e4c43421efbcff799 Mon Sep 17 00:00:00 2001 From: yucheng Date: Thu, 10 Sep 2026 07:32:54 +0000 Subject: [PATCH 001/112] fix(otel): cap per-index OpenInference message attributes span-wide OpenInferenceMapper spelled every captured prompt and response message out as two indexed attributes with no bound. A few dozen turns overran the OTel SDK's 128-attribute span limit, which evicts oldest first, so the gen_ai.* model, provider, usage, cost and finish reason written before it were what got dropped. Both directions now share one MAX_MESSAGE_ATTRS_PER_SPAN ceiling, the response keeps at least half of it, and input.value / output.value still carry the complete conversation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../otel/mappers/openinference.py | 26 ++- litellm/integrations/otel/mappers/utils.py | 12 ++ .../integrations/otel/test_otel_v2_emitter.py | 148 ++++++++++++++++++ 3 files changed, 181 insertions(+), 5 deletions(-) diff --git a/litellm/integrations/otel/mappers/openinference.py b/litellm/integrations/otel/mappers/openinference.py index 0ba45170b8e..1e2dbf6974d 100644 --- a/litellm/integrations/otel/mappers/openinference.py +++ b/litellm/integrations/otel/mappers/openinference.py @@ -12,6 +12,7 @@ from typing import Final from litellm.integrations.otel.mappers.base import AttributeMap, AttrValue, SpanData from litellm.integrations.otel.mappers.utils import ( + MAX_MESSAGE_ATTRS_PER_SPAN, MAX_TOOL_DEFINITION_ATTRS_PER_SPAN, collect, drop_none, @@ -26,6 +27,8 @@ from litellm.integrations.otel.model.payloads import ( ToolDefinition, ) +_MAX_INDEXED_MESSAGES: Final = MAX_MESSAGE_ATTRS_PER_SPAN // 2 + class OpenInferenceMapper: """Emits OpenInference attributes for LLM_CALL spans. @@ -84,22 +87,35 @@ class OpenInferenceMapper: return {} def _llm_call(self, data: LLMCallSpanData) -> AttributeMap: + outputs: Final = output_messages(data) + indexed_in, indexed_out = self._indexed_split(len(data.messages_in), len(outputs)) return { **collect(self._LLM_CALL_ATTRS, data), **collect(self._BLOB_ATTRS, data), - **self._messages("llm.input_messages", "input.value", data.messages_in), - **self._messages("llm.output_messages", "output.value", output_messages(data)), + **self._messages("llm.input_messages", "input.value", data.messages_in, indexed_in), + **self._messages("llm.output_messages", "output.value", outputs, indexed_out), **self._tools(data), } @staticmethod - def _messages(prefix: str, value_key: str, messages: Sequence[object]) -> AttributeMap: - """Per-message ``{prefix}.{idx}.message.*`` keys + the ``value_key`` blob.""" + def _indexed_split(inputs: int, outputs: int) -> tuple[int, int]: + """How many prompt and response messages get per-index attributes. + + Both directions share one span-wide allowance. The response is reserved at + least half of it, so a long prompt can never push the completion off the + span, and the prompt takes whatever the response leaves unused. + """ + indexed_out: Final = min(outputs, max(_MAX_INDEXED_MESSAGES // 2, _MAX_INDEXED_MESSAGES - inputs)) + return _MAX_INDEXED_MESSAGES - indexed_out, indexed_out + + @staticmethod + def _messages(prefix: str, value_key: str, messages: Sequence[object], indexed: int) -> AttributeMap: + """``{prefix}.{idx}.message.*`` keys for the leading ``indexed`` messages + the ``value_key`` blob of all.""" parsed: Final = [(m.get("role") if isinstance(m, dict) else None, message_content(m)) for m in messages] attrs: Final = drop_none( { key: value - for idx, (role, content) in enumerate(parsed) + for idx, (role, content) in enumerate(parsed[:indexed]) for key, value in ( ( f"{prefix}.{idx}.message.role", diff --git a/litellm/integrations/otel/mappers/utils.py b/litellm/integrations/otel/mappers/utils.py index d45dca782b2..8d21b774319 100644 --- a/litellm/integrations/otel/mappers/utils.py +++ b/litellm/integrations/otel/mappers/utils.py @@ -32,6 +32,18 @@ core telemetry no matter how many vocabularies are configured. """ +MAX_MESSAGE_ATTRS_PER_SPAN: Final = DEFAULT_SPAN_ATTRIBUTE_LIMIT // 4 +"""Span-wide ceiling on attributes spent spelling out chat messages per index. + +A conversation is the other unbounded family: two attributes per message, for +the prompt and the response alike, on the same span. Past a few dozen turns the +family alone exceeds the span attribute limit and evicts the core telemetry +written before it. The ceiling covers both directions together, since a budget +handed to each direction separately doubles. The complete conversation still +rides the JSON blob attributes; only the per-index convenience keys are capped. +""" + + def tool_attr_budget(vocabularies: int) -> int: """Split the span-wide tool-definition ceiling across active vocabularies.""" return MAX_TOOL_DEFINITION_ATTRS_PER_SPAN // max(vocabularies, 1) diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py b/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py index b1b1b62c820..a417bd62124 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py @@ -1,6 +1,8 @@ """Golden tests for the OTel v2 engine: span shape, kinds, semconv attributes, legacy dual-emit, hierarchy, error status, and idempotency. Needs the OTel SDK.""" +import json + import pytest pytest.importorskip("opentelemetry") @@ -18,6 +20,7 @@ from litellm.integrations.otel.plumbing import providers # noqa: E402 from litellm.integrations.otel.emitter import SpanEmitter # noqa: E402 from litellm.integrations.otel.emitter import stamp_error # noqa: E402 from litellm.integrations.otel.mappers.utils import ( # noqa: E402 + MAX_MESSAGE_ATTRS_PER_SPAN, MAX_TOOL_DEFINITION_ATTRS_PER_SPAN, ) from litellm.integrations.otel.model.payloads import ( # noqa: E402 @@ -440,3 +443,148 @@ def test_vendor_tool_definitions_are_truncated_not_dropped(): assert a["llm.tools.0.tool.name"] == "tool_0" assert a["llm.tools.0.tool.json_schema"] assert "llm.tools.126.tool.name" not in a + + +def _conversation_payload(turns, choices=1, **overrides): + """A ``turns``-message chat with ``choices`` response choices, content-bearing.""" + return _payload( + messages=[{"role": ("user", "assistant")[i % 2], "content": f"turn {i}"} for i in range(turns)], + response={ + "id": "resp_1", + "model": "gpt-4o-2024", + "choices": [ + {"finish_reason": "stop", "message": {"role": "assistant", "content": f"reply {i}"}} + for i in range(choices) + ], + }, + **overrides, + ) + + +def _conversation_span(mapper_names, payload): + """The exported LLM-call span for ``payload`` with content capture on.""" + cfg = OpenTelemetryV2Config( + exporter="in_memory", + mapper_names=list(mapper_names), + capture_message_content="span_only", + ) + provider, exporter = providers.in_memory_provider(cfg) + engine = SpanEmitter(providers.get_tracer(provider, "litellm-test"), cfg) + engine.emit( + SpanRole.LLM_CALL, + LLMCallSpanData.from_standard_logging_payload(payload, capture_content=True), + ) + (span,) = exporter.get_finished_spans() + return span + + +def _indexed_message_count(attributes, prefix): + return len({key.split(".")[2] for key in attributes if key.startswith(f"{prefix}.")}) + + +@pytest.mark.parametrize("turns", [60, 200]) +def test_long_conversation_does_not_evict_core_attributes(turns): + """Per-message OpenInference attributes must never crowd core telemetry off the span. + + With content capture on, the OpenInference vocabulary spells every prompt and + response message out as two per-index attributes. A few dozen turns overruns + the OTel SDK's 128-attribute span limit, which evicts oldest-first, so the + ``gen_ai.*`` set written before it is what disappears. + """ + span = _conversation_span(["genai", "openinference"], _conversation_payload(turns)) + a = span.attributes + + assert span.dropped_attributes == 0 + assert a[GenAI.REQUEST_MODEL] == "gpt-4o" + assert a[GenAI.PROVIDER_NAME] == "openai" + assert a[GenAI.USAGE_INPUT_TOKENS] == 10 + assert a[GenAI.USAGE_OUTPUT_TOKENS] == 5 + assert a[GenAI.RESPONSE_FINISH_REASONS] == ("stop",) + assert a[f"{LiteLLM.COST_PREFIX}total"] == 0.002 + + assert a["llm.input_messages.0.message.content"] == "turn 0" + assert a["llm.output_messages.0.message.content"] == "reply 0" + assert f"llm.input_messages.{turns - 1}.message.role" not in a + assert len(json.loads(a["input.value"])) == turns + assert len(json.loads(a["output.value"])) == 1 + assert len(json.loads(a[GenAI.INPUT_MESSAGES])) == turns + + +def test_short_conversation_keeps_every_message_indexed(): + """Below the cap nothing is truncated in either direction.""" + a = _conversation_span(["genai", "openinference"], _conversation_payload(4, choices=2)).attributes + for idx in range(4): + assert a[f"llm.input_messages.{idx}.message.content"] == f"turn {idx}" + for idx in range(2): + assert a[f"llm.output_messages.{idx}.message.content"] == f"reply {idx}" + + +def test_message_cap_is_shared_across_input_and_output(): + """One span-wide allowance covers both directions, and the response always keeps a share. + + A long prompt takes what a single reply leaves over, and a many-choice reply + cannot take the whole allowance away from the prompt either. + """ + long_prompt = _conversation_span(["genai", "openinference"], _conversation_payload(60, choices=1)).attributes + many_choices = _conversation_span(["genai", "openinference"], _conversation_payload(60, choices=20)).attributes + + single_reply_indexed = _indexed_message_count(long_prompt, "llm.output_messages") + assert single_reply_indexed == 1 + assert _indexed_message_count(long_prompt, "llm.input_messages") + single_reply_indexed == ( + MAX_MESSAGE_ATTRS_PER_SPAN // 2 + ) + + assert _indexed_message_count(many_choices, "llm.input_messages") > 0 + assert _indexed_message_count(many_choices, "llm.output_messages") > single_reply_indexed + assert _indexed_message_count(many_choices, "llm.input_messages") + _indexed_message_count( + many_choices, "llm.output_messages" + ) == (MAX_MESSAGE_ATTRS_PER_SPAN // 2) + + +def test_fully_populated_arize_span_stays_within_the_attribute_limit(): + """Every capped family maxed at once still leaves the whole core intact. + + The Arize / Phoenix composition (``genai`` + ``openinference`` + ``legacy``) + with every request parameter, every cost component, a hundred-plus tools, a + two-hundred-turn prompt and twenty choices is the worst case the two + span-wide ceilings have to absorb together. + """ + payload = _conversation_payload( + 200, + choices=20, + stream=True, + model_parameters={ + **_tools_payload(127)["model_parameters"], + "top_p": 0.9, + "frequency_penalty": 0.1, + "presence_penalty": 0.1, + "seed": 7, + "stop": ["\n"], + }, + cost_breakdown={ + key: 0.001 + for key in ( + "input_cost", + "output_cost", + "cache_read_cost", + "cache_creation_cost", + "tool_usage_cost", + "original_cost", + "discount_amount", + "discount_percent", + "margin_fixed_amount", + "margin_percent", + "margin_total_amount", + "total_cost", + ) + }, + ) + span = _conversation_span(["genai", "openinference"], payload) + a = span.attributes + + assert span.dropped_attributes == 0 + assert a[GenAI.REQUEST_MODEL] == "gpt-4o" + assert a[f"{LiteLLM.COST_PREFIX}total"] == 0.002 + assert a[LiteLLM.TOOLS_DECLARED] == 127 + assert a["llm.input_messages.0.message.content"] == "turn 0" + assert a["llm.output_messages.0.message.content"] == "reply 0" From fcaf2d7d98164fc8561c411a999ffd42469797e9 Mon Sep 17 00:00:00 2001 From: yucheng Date: Thu, 10 Sep 2026 08:00:51 +0000 Subject: [PATCH 002/112] fix(otel): size the message ceiling so every vocabulary fits beside it Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/otel/mappers/utils.py | 9 ++++++--- .../integrations/otel/test_otel_v2_emitter.py | 15 ++++++++------- 2 files changed, 14 insertions(+), 10 deletions(-) diff --git a/litellm/integrations/otel/mappers/utils.py b/litellm/integrations/otel/mappers/utils.py index 8d21b774319..d8918491720 100644 --- a/litellm/integrations/otel/mappers/utils.py +++ b/litellm/integrations/otel/mappers/utils.py @@ -32,15 +32,18 @@ core telemetry no matter how many vocabularies are configured. """ -MAX_MESSAGE_ATTRS_PER_SPAN: Final = DEFAULT_SPAN_ATTRIBUTE_LIMIT // 4 +MAX_MESSAGE_ATTRS_PER_SPAN: Final = DEFAULT_SPAN_ATTRIBUTE_LIMIT // 8 """Span-wide ceiling on attributes spent spelling out chat messages per index. A conversation is the other unbounded family: two attributes per message, for the prompt and the response alike, on the same span. Past a few dozen turns the family alone exceeds the span attribute limit and evicts the core telemetry written before it. The ceiling covers both directions together, since a budget -handed to each direction separately doubles. The complete conversation still -rides the JSON blob attributes; only the per-index convenience keys are capped. +handed to each direction separately doubles. An eighth is the largest share +that still fits beside the tool ceiling and the core of every vocabulary at +once, request parameters, cost breakdown and identity included. The complete +conversation still rides the JSON blob attributes; only the per-index +convenience keys are capped. """ diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py b/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py index a417bd62124..f571ab7004b 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py @@ -461,10 +461,11 @@ def _conversation_payload(turns, choices=1, **overrides): ) -def _conversation_span(mapper_names, payload): +def _conversation_span(mapper_names, payload, legacy_compat=False): """The exported LLM-call span for ``payload`` with content capture on.""" cfg = OpenTelemetryV2Config( exporter="in_memory", + legacy_compat=legacy_compat, mapper_names=list(mapper_names), capture_message_content="span_only", ) @@ -541,13 +542,13 @@ def test_message_cap_is_shared_across_input_and_output(): ) == (MAX_MESSAGE_ATTRS_PER_SPAN // 2) -def test_fully_populated_arize_span_stays_within_the_attribute_limit(): +def test_fully_populated_span_with_every_vocabulary_stays_within_the_attribute_limit(): """Every capped family maxed at once still leaves the whole core intact. - The Arize / Phoenix composition (``genai`` + ``openinference`` + ``legacy``) - with every request parameter, every cost component, a hundred-plus tools, a - two-hundred-turn prompt and twenty choices is the worst case the two - span-wide ceilings have to absorb together. + Every vocabulary in the registry plus ``legacy``, every request parameter, + every cost component, a hundred-plus tools, a two-hundred-turn prompt and + twenty choices is the worst case the two span-wide ceilings have to absorb + together. """ payload = _conversation_payload( 200, @@ -579,7 +580,7 @@ def test_fully_populated_arize_span_stays_within_the_attribute_limit(): ) }, ) - span = _conversation_span(["genai", "openinference"], payload) + span = _conversation_span(["genai", "openinference", "langfuse", "weave", "langtrace"], payload, legacy_compat=True) a = span.attributes assert span.dropped_attributes == 0 From 302a8d43da054f98b7ccb75baf43a8a79bb3ea40 Mon Sep 17 00:00:00 2001 From: shivam Date: Thu, 10 Sep 2026 22:02:58 +0000 Subject: [PATCH 003/112] fix(cost): bill cached realtime audio tokens at the audio cache-read rate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/cost_calculator.py | 78 +++++++++++++------ .../litellm_core_utils/llm_cost_calc/utils.py | 40 ++++++++-- ...odel_prices_and_context_window_backup.json | 3 + .../transformation.py | 45 ++++++----- litellm/responses/utils.py | 3 + litellm/types/llms/openai.py | 14 ++++ litellm/types/utils.py | 6 ++ model_prices_and_context_window.json | 3 + .../llm_cost_calc/test_llm_cost_calc_utils.py | 68 ++++++++++++++++ .../responses/test_responses_utils.py | 41 ++++++++++ tests/test_litellm/test_cost_calculator.py | 61 +++++++++++++++ 11 files changed, 314 insertions(+), 48 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 814eaaf76f7..7cd3ea8f303 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -108,6 +108,7 @@ from litellm.types.llms.openai import ( ) from litellm.types.rerank import RerankBilledUnits, RerankResponse from litellm.types.utils import ( + CachedTokensDetails, CallTypesLiteral, LiteLLMRealtimeStreamLoggingObject, LlmProviders, @@ -2310,6 +2311,60 @@ def _summable_prompt_token_fields(prompt_tokens_details: BaseModel) -> list[str] return [attr for attr in field_names if attr != "cache_creation_tokens"] +def _combine_cached_tokens_details( + current: CachedTokensDetails | None, new: CachedTokensDetails +) -> CachedTokensDetails: + def _sum_optional(current_value: int | None, new_value: int | None) -> int | None: + if current_value is None and new_value is None: + return None + return (current_value or 0) + (new_value or 0) + + return CachedTokensDetails( + text_tokens=_sum_optional( + current.text_tokens if current is not None else None, new.text_tokens + ), + audio_tokens=_sum_optional( + current.audio_tokens if current is not None else None, new.audio_tokens + ), + image_tokens=_sum_optional( + current.image_tokens if current is not None else None, new.image_tokens + ), + ) + + +def _combine_prompt_tokens_details(combined: Usage, usage: Usage) -> None: + if not (hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details): + return + if not hasattr(combined, "prompt_tokens_details") or not combined.prompt_tokens_details: + combined.prompt_tokens_details = PromptTokensDetailsWrapper() + + # Check what keys exist in the model's prompt_tokens_details + # Access model_fields on the class, not the instance, to avoid Pydantic 2.11+ deprecation warnings + for attr in _summable_prompt_token_fields(usage.prompt_tokens_details): + if ( + hasattr(usage.prompt_tokens_details, attr) + and not attr.startswith("_") + and not callable(_attribute_value(usage.prompt_tokens_details, attr)) + ): + current_val = getattr(combined.prompt_tokens_details, attr, 0) or 0 + new_val = getattr(usage.prompt_tokens_details, attr, 0) or 0 + if new_val is not None and isinstance(new_val, (int, float)): + setattr( + combined.prompt_tokens_details, + attr, + current_val + new_val, + ) + + new_cached_tokens_details: Final = getattr( + usage.prompt_tokens_details, "cached_tokens_details", None + ) + if isinstance(new_cached_tokens_details, CachedTokensDetails): + combined.prompt_tokens_details.cached_tokens_details = _combine_cached_tokens_details( + getattr(combined.prompt_tokens_details, "cached_tokens_details", None), + new_cached_tokens_details, + ) + + class BaseTokenUsageProcessor: @staticmethod def combine_usage_objects(usage_objects: list[Usage]) -> Usage: @@ -2318,7 +2373,6 @@ class BaseTokenUsageProcessor: """ from litellm.types.utils import ( CompletionTokensDetailsWrapper, - PromptTokensDetailsWrapper, Usage, ) @@ -2337,27 +2391,7 @@ class BaseTokenUsageProcessor: and isinstance(current_val, (int, float)) ): setattr(combined, attr, current_val + new_val) - # Handle nested prompt_tokens_details - if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details: - if not hasattr(combined, "prompt_tokens_details") or not combined.prompt_tokens_details: - combined.prompt_tokens_details = PromptTokensDetailsWrapper() - - # Check what keys exist in the model's prompt_tokens_details - # Access model_fields on the class, not the instance, to avoid Pydantic 2.11+ deprecation warnings - for attr in _summable_prompt_token_fields(usage.prompt_tokens_details): - if ( - hasattr(usage.prompt_tokens_details, attr) - and not attr.startswith("_") - and not callable(_attribute_value(usage.prompt_tokens_details, attr)) - ): - current_val = getattr(combined.prompt_tokens_details, attr, 0) or 0 - new_val = getattr(usage.prompt_tokens_details, attr, 0) or 0 - if new_val is not None and isinstance(new_val, (int, float)): - setattr( - combined.prompt_tokens_details, - attr, - current_val + new_val, - ) + _combine_prompt_tokens_details(combined, usage) # Handle nested completion_tokens_details if hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details: diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index e5977ca4156..18ef99597a0 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -9,6 +9,8 @@ from types import MappingProxyType from typing import Any, Final, Literal, TypedDict, cast from zoneinfo import ZoneInfo, ZoneInfoNotFoundError +from typing_extensions import ReadOnly + import litellm from litellm._internal_context import current_billing_time from litellm._logging import verbose_logger @@ -772,6 +774,7 @@ def calculate_cache_writing_cost( class PromptTokensDetailsResult(TypedDict): cache_hit_tokens: int + cache_hit_audio_tokens: ReadOnly[int] cache_creation_tokens: int cache_creation_token_details: CacheCreationTokenDetails | None text_tokens: int @@ -802,12 +805,26 @@ def parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: ) or None ) - text_tokens: Final = ( - cast(int | None, getattr(usage.prompt_tokens_details, "text_tokens", None)) - or 0 # default to prompt tokens, if this field is not set + cached_tokens_details: Final = getattr(usage.prompt_tokens_details, "cached_tokens_details", None) + cached_text_tokens: Final = _get_token_detail_value(cached_tokens_details, "text_tokens") or 0 + cached_audio_tokens: Final = _get_token_detail_value(cached_tokens_details, "audio_tokens") or 0 + cached_image_tokens: Final = _get_token_detail_value(cached_tokens_details, "image_tokens") or 0 + text_tokens: Final = max( + ( + cast(int | None, getattr(usage.prompt_tokens_details, "text_tokens", None)) + or 0 # default to prompt tokens, if this field is not set + ) + - cached_text_tokens, + 0, + ) + audio_tokens: Final = max( + (cast(int | None, getattr(usage.prompt_tokens_details, "audio_tokens", 0)) or 0) - cached_audio_tokens, + 0, + ) + image_tokens: Final = max( + (cast(int | None, getattr(usage.prompt_tokens_details, "image_tokens", 0)) or 0) - cached_image_tokens, + 0, ) - audio_tokens: Final = cast(int | None, getattr(usage.prompt_tokens_details, "audio_tokens", 0)) or 0 - image_tokens: Final = cast(int | None, getattr(usage.prompt_tokens_details, "image_tokens", 0)) or 0 video_tokens: Final = _coerce_token_count(getattr(usage.prompt_tokens_details, "video_tokens", 0)) character_count: Final = ( cast( @@ -835,6 +852,7 @@ def parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: return PromptTokensDetailsResult( cache_hit_tokens=cache_hit_tokens, + cache_hit_audio_tokens=min(cached_audio_tokens, cache_hit_tokens), cache_creation_tokens=cache_creation_tokens, cache_creation_token_details=cache_creation_token_details, text_tokens=text_tokens, @@ -918,7 +936,16 @@ def _calculate_input_cost( prompt_cost = float(prompt_tokens_details["text_tokens"]) * prompt_base_cost ### CACHE READ COST - Now uses tiered pricing - prompt_cost += float(prompt_tokens_details["cache_hit_tokens"]) * cache_read_cost + cache_hit_audio_tokens: Final = prompt_tokens_details["cache_hit_audio_tokens"] + audio_cache_read_rate: Final = _get_cost_per_unit( + model_info, + _get_service_tier_cost_key("cache_read_input_audio_token_cost", service_tier), + None, + ) + prompt_cost += float(prompt_tokens_details["cache_hit_tokens"] - cache_hit_audio_tokens) * cache_read_cost + prompt_cost += float(cache_hit_audio_tokens) * ( + audio_cache_read_rate if audio_cache_read_rate is not None else cache_read_cost + ) ### AUDIO COST if prompt_tokens_details["audio_tokens"]: @@ -1149,6 +1176,7 @@ def generic_cost_per_token( ### PROCESSING COST prompt_tokens_details = PromptTokensDetailsResult( cache_hit_tokens=0, + cache_hit_audio_tokens=0, cache_creation_tokens=0, cache_creation_token_details=None, text_tokens=usage.prompt_tokens, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 1c0bd32d782..7f4ca991bd3 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -32502,6 +32502,7 @@ }, "gpt-realtime": { "cache_creation_input_audio_token_cost": 4e-07, + "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 3.2e-05, @@ -32535,6 +32536,7 @@ }, "gpt-realtime-1.5": { "cache_creation_input_audio_token_cost": 4e-07, + "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, "input_cost_per_audio_token": 3.2e-05, "input_cost_per_image_token": 5e-06, @@ -32702,6 +32704,7 @@ }, "gpt-realtime-2025-08-28": { "cache_creation_input_audio_token_cost": 4e-07, + "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 3.2e-05, diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index fca5b0d11cf..cc84c68b0f5 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -45,6 +45,7 @@ from litellm.responses.litellm_completion_transformation.session_handler import ) from litellm.types.llms.openai import ( AllMessageValues, + CachedTokensDetails, ChatCompletionImageObject, ChatCompletionImageUrlObject, ChatCompletionRedactedThinkingBlock, @@ -2681,27 +2682,31 @@ class LiteLLMCompletionResponsesConfig: # Translate prompt_tokens_details to input_tokens_details if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details is not None: prompt_details: Final = usage.prompt_tokens_details - input_details_dict: Final[dict[str, int]] = {} - - if hasattr(prompt_details, "cached_tokens") and prompt_details.cached_tokens is not None: - input_details_dict["cached_tokens"] = prompt_details.cached_tokens - else: - input_details_dict["cached_tokens"] = 0 - - if hasattr(prompt_details, "text_tokens") and prompt_details.text_tokens is not None: - input_details_dict["text_tokens"] = prompt_details.text_tokens - - if hasattr(prompt_details, "audio_tokens") and prompt_details.audio_tokens is not None: - input_details_dict["audio_tokens"] = prompt_details.audio_tokens - - cache_write_tokens = getattr(prompt_details, "cache_write_tokens", None) or getattr( - prompt_details, "cache_creation_tokens", None + cached_tokens_details: Final = getattr(prompt_details, "cached_tokens_details", None) + response_usage.input_tokens_details = InputTokensDetails( + cached_tokens=( + prompt_details.cached_tokens + if hasattr(prompt_details, "cached_tokens") and prompt_details.cached_tokens is not None + else 0 + ), + text_tokens=( + prompt_details.text_tokens + if hasattr(prompt_details, "text_tokens") and prompt_details.text_tokens is not None + else None + ), + audio_tokens=( + prompt_details.audio_tokens + if hasattr(prompt_details, "audio_tokens") and prompt_details.audio_tokens is not None + else None + ), + cache_write_tokens=( + getattr(prompt_details, "cache_write_tokens", None) + or getattr(prompt_details, "cache_creation_tokens", None) + ), + cached_tokens_details=( + cached_tokens_details if isinstance(cached_tokens_details, CachedTokensDetails) else None + ), ) - if cache_write_tokens is not None: - input_details_dict["cache_write_tokens"] = cache_write_tokens - - if input_details_dict: - response_usage.input_tokens_details = InputTokensDetails(**input_details_dict) # Translate completion_tokens_details to output_tokens_details if hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details is not None: diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 599e978df6a..d63e3ddf0aa 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -1179,6 +1179,9 @@ class ResponseAPILoggingUtils: audio_tokens=getattr(response_api_usage.input_tokens_details, "audio_tokens", None), text_tokens=getattr(response_api_usage.input_tokens_details, "text_tokens", None), image_tokens=getattr(response_api_usage.input_tokens_details, "image_tokens", None), + cached_tokens_details=getattr( + response_api_usage.input_tokens_details, "cached_tokens_details", None + ), cache_write_tokens=getattr(response_api_usage.input_tokens_details, "cache_write_tokens", None), ) completion_tokens_details: CompletionTokensDetailsWrapper | None = None diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index b7c4371f32f..274747b4193 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -1284,9 +1284,16 @@ class OutputTokensDetails(BaseLiteLLMOpenAIResponseObject): model_config = {"extra": "allow"} +class CachedTokensDetails(BaseModel): + text_tokens: int | None = None + audio_tokens: int | None = None + image_tokens: int | None = None + + class InputTokensDetails(BaseLiteLLMOpenAIResponseObject): audio_tokens: int | None = None cached_tokens: int = 0 + cached_tokens_details: CachedTokensDetails | None = None text_tokens: int | None = None model_config = {"extra": "allow"} @@ -2204,10 +2211,17 @@ class OpenAIRealtimeInputAudioTranscriptionCompleted(TypedDict): transcript: ReadOnly[str] +class OpenAIRealtimeCachedTokensDetails(TypedDict, total=False): + text_tokens: ReadOnly[int] + audio_tokens: ReadOnly[int] + image_tokens: ReadOnly[int] + + class OpenAIRealtimeUsageTokenDetails(TypedDict): audio_tokens: ReadOnly[int] text_tokens: ReadOnly[int] cached_tokens: NotRequired[ReadOnly[int]] + cached_tokens_details: NotRequired[ReadOnly[OpenAIRealtimeCachedTokensDetails]] class OpenAIRealtimeResponseUsage(TypedDict): diff --git a/litellm/types/utils.py b/litellm/types/utils.py index ab0cc5f959c..39100031dcf 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -58,6 +58,7 @@ from .llms.base import HiddenParams from .llms.openai import ( AllMessageValues, Batch, + CachedTokensDetails, ChatCompletionAnnotation, ChatCompletionReasoningItem, ChatCompletionRedactedThinkingBlock, @@ -1707,6 +1708,9 @@ class PromptTokensDetailsWrapper( cache_creation_token_details: CacheCreationTokenDetails | None = None """Details of cache creation tokens sent to the model. Used for tracking 5m/1h cache creation tokens for Anthropic prompt caching.""" + cached_tokens_details: CachedTokensDetails | None = None + """Details of cached (cache-hit) tokens sent to the model. OpenAI realtime naming; carries the per-modality cache-read split.""" + def __setattr__(self, name: str, value: object) -> None: super().__setattr__(name, value) if name == "cache_write_tokens": @@ -1753,6 +1757,8 @@ class PromptTokensDetailsWrapper( del self.cache_creation_tokens if self.cache_creation_token_details is None: del self.cache_creation_token_details + if self.cached_tokens_details is None: + del self.cached_tokens_details class ServerToolUse(BaseModel): diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 1c0bd32d782..7f4ca991bd3 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -32502,6 +32502,7 @@ }, "gpt-realtime": { "cache_creation_input_audio_token_cost": 4e-07, + "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 3.2e-05, @@ -32535,6 +32536,7 @@ }, "gpt-realtime-1.5": { "cache_creation_input_audio_token_cost": 4e-07, + "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, "input_cost_per_audio_token": 3.2e-05, "input_cost_per_image_token": 5e-06, @@ -32702,6 +32704,7 @@ }, "gpt-realtime-2025-08-28": { "cache_creation_input_audio_token_cost": 4e-07, + "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 3.2e-05, diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index fbb9d178390..3709b526c3b 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -2648,6 +2648,7 @@ def test_cache_writing_cost_with_zero_creation_tokens_and_ephemeral_details(): prompt_tokens_details: PromptTokensDetailsResult = { "cache_hit_tokens": 0, + "cache_hit_audio_tokens": 0, "cache_creation_tokens": 0, "cache_creation_token_details": CacheCreationTokenDetails( ephemeral_5m_input_tokens=100, @@ -5147,3 +5148,70 @@ def test_generic_cost_per_token_bills_nested_reasoning_once_beside_audio_output( assert completion_cost == pytest.approx( 30 * info["output_cost_per_token"] + 70 * info["output_cost_per_audio_token"] ) + + +def test_cached_realtime_audio_tokens_billed_at_audio_cache_read_rate( + _local_model_cost_map: None, +) -> None: + usage = Usage( + prompt_tokens=283, + completion_tokens=0, + total_tokens=283, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=116, + audio_tokens=167, + cached_tokens=192, + cached_tokens_details={"text_tokens": 64, "audio_tokens": 128}, + ), + ) + + prompt_cost, _ = generic_cost_per_token( + model="gpt-realtime-2", usage=usage, custom_llm_provider="openai" + ) + assert prompt_cost == pytest.approx(0.0015328) + + +def test_prompt_tokens_details_without_cached_tokens_details_unchanged( + _local_model_cost_map: None, +) -> None: + usage = Usage( + prompt_tokens=283, + completion_tokens=0, + total_tokens=283, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=116, audio_tokens=167, cached_tokens=192 + ), + ) + + prompt_cost, _ = generic_cost_per_token( + model="gpt-realtime-2", usage=usage, custom_llm_provider="openai" + ) + assert prompt_cost == pytest.approx(0.0029888) + + +def test_cached_audio_tokens_fall_back_to_cache_read_input_token_cost() -> None: + model_info: ModelInfo = { + "input_cost_per_token": 4e-6, + "input_cost_per_audio_token": 32e-6, + "cache_read_input_token_cost": 5e-7, + } + usage = Usage( + prompt_tokens=283, + completion_tokens=0, + total_tokens=283, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=116, + audio_tokens=167, + cached_tokens=192, + cached_tokens_details={"text_tokens": 64, "audio_tokens": 128}, + ), + ) + + prompt_cost, _ = generic_cost_per_token( + model="some-realtime-model", + usage=usage, + custom_llm_provider="openai", + model_info=model_info, + ) + expected = 52 * 4e-6 + 64 * 5e-7 + 39 * 32e-6 + 128 * 5e-7 + assert prompt_cost == pytest.approx(expected) diff --git a/tests/test_litellm/responses/test_responses_utils.py b/tests/test_litellm/responses/test_responses_utils.py index 9d9eefdceb3..4d06b5e7bdc 100644 --- a/tests/test_litellm/responses/test_responses_utils.py +++ b/tests/test_litellm/responses/test_responses_utils.py @@ -577,6 +577,47 @@ class TestResponseAPILoggingUtils: assert result.completion_tokens_details is not None assert result.completion_tokens_details.reasoning_tokens == 4 + def test_transform_realtime_usage_dict_keeps_cached_tokens_details(self): + usage = { + "input_tokens": 283, + "output_tokens": 0, + "total_tokens": 283, + "input_token_details": { + "text_tokens": 116, + "audio_tokens": 167, + "cached_tokens": 192, + "cached_tokens_details": {"text_tokens": 64, "audio_tokens": 128}, + }, + } + + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) + + assert result.prompt_tokens_details is not None + assert result.prompt_tokens_details.cached_tokens == 192 + assert result.prompt_tokens_details.cached_tokens_details is not None + assert result.prompt_tokens_details.cached_tokens_details.audio_tokens == 128 + assert result.prompt_tokens_details.cached_tokens_details.text_tokens == 64 + + def test_transform_response_api_usage_object_keeps_cached_tokens_details(self): + usage = ResponseAPIUsage( + input_tokens=283, + output_tokens=0, + total_tokens=283, + input_tokens_details={ + "text_tokens": 116, + "audio_tokens": 167, + "cached_tokens": 192, + "cached_tokens_details": {"text_tokens": 64, "audio_tokens": 128}, + }, + ) + + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) + + assert result.prompt_tokens_details is not None + assert result.prompt_tokens_details.cached_tokens_details is not None + assert result.prompt_tokens_details.cached_tokens_details.audio_tokens == 128 + assert result.prompt_tokens_details.cached_tokens_details.text_tokens == 64 + class TestResponsesAPIProviderSpecificParams: """ diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index f610821e06a..0b339e3d525 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -4768,3 +4768,64 @@ def test_collect_and_combine_realtime_usage_stores_partitioned_text_tokens() -> assert combined.completion_tokens_details.reasoning_tokens == 95 assert combined.completion_tokens_details.text_tokens == 38 assert combined.completion_tokens_details.audio_tokens == 0 + + +def test_realtime_combine_sums_nested_cached_tokens_details(): + results: OpenAIRealtimeStreamList = [ + { + "type": "response.done", + "response": { + "usage": { + "input_tokens": 283, + "output_tokens": 0, + "total_tokens": 283, + "input_token_details": { + "text_tokens": 116, + "audio_tokens": 167, + "cached_tokens": 192, + "cached_tokens_details": {"text_tokens": 64, "audio_tokens": 128}, + }, + } + }, + }, + { + "type": "response.done", + "response": { + "usage": { + "input_tokens": 150, + "output_tokens": 0, + "total_tokens": 150, + "input_token_details": { + "text_tokens": 50, + "audio_tokens": 100, + "cached_tokens": 100, + "cached_tokens_details": {"audio_tokens": 100}, + }, + } + }, + }, + ] + + combined = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results( + results=results, + ) + + assert combined.prompt_tokens_details is not None + assert combined.prompt_tokens_details.cached_tokens == 292 + assert combined.prompt_tokens_details.cached_tokens_details is not None + assert combined.prompt_tokens_details.cached_tokens_details.audio_tokens == 228 + assert combined.prompt_tokens_details.cached_tokens_details.text_tokens == 64 + assert combined.prompt_tokens_details.cached_tokens_details.image_tokens is None + + +def test_usage_without_cached_tokens_details_omits_key(): + usage = Usage( + prompt_tokens=10, + completion_tokens=5, + total_tokens=15, + prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=10), + ) + + dumped = usage.prompt_tokens_details.model_dump() + assert "cached_tokens_details" not in dumped + assert "cached_tokens_details" not in usage.prompt_tokens_details.model_dump_json() From 67fc9e4e3dcb94035c4b4d07d63c3565939d9a3f Mon Sep 17 00:00:00 2001 From: shivam Date: Thu, 10 Sep 2026 22:10:06 +0000 Subject: [PATCH 004/112] fix(responses): only emit cache_write_tokens when reported Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/cost_calculator.py | 16 +++------- .../transformation.py | 30 +++++++------------ .../test_litellm_completion_responses.py | 2 ++ 3 files changed, 16 insertions(+), 32 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 7cd3ea8f303..8daa2de416b 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2320,15 +2320,9 @@ def _combine_cached_tokens_details( return (current_value or 0) + (new_value or 0) return CachedTokensDetails( - text_tokens=_sum_optional( - current.text_tokens if current is not None else None, new.text_tokens - ), - audio_tokens=_sum_optional( - current.audio_tokens if current is not None else None, new.audio_tokens - ), - image_tokens=_sum_optional( - current.image_tokens if current is not None else None, new.image_tokens - ), + text_tokens=_sum_optional(current.text_tokens if current is not None else None, new.text_tokens), + audio_tokens=_sum_optional(current.audio_tokens if current is not None else None, new.audio_tokens), + image_tokens=_sum_optional(current.image_tokens if current is not None else None, new.image_tokens), ) @@ -2355,9 +2349,7 @@ def _combine_prompt_tokens_details(combined: Usage, usage: Usage) -> None: current_val + new_val, ) - new_cached_tokens_details: Final = getattr( - usage.prompt_tokens_details, "cached_tokens_details", None - ) + new_cached_tokens_details: Final = getattr(usage.prompt_tokens_details, "cached_tokens_details", None) if isinstance(new_cached_tokens_details, CachedTokensDetails): combined.prompt_tokens_details.cached_tokens_details = _combine_cached_tokens_details( getattr(combined.prompt_tokens_details, "cached_tokens_details", None), diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index cc84c68b0f5..e6f90b99b60 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -2683,30 +2683,20 @@ class LiteLLMCompletionResponsesConfig: if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details is not None: prompt_details: Final = usage.prompt_tokens_details cached_tokens_details: Final = getattr(prompt_details, "cached_tokens_details", None) - response_usage.input_tokens_details = InputTokensDetails( - cached_tokens=( - prompt_details.cached_tokens - if hasattr(prompt_details, "cached_tokens") and prompt_details.cached_tokens is not None - else 0 - ), - text_tokens=( - prompt_details.text_tokens - if hasattr(prompt_details, "text_tokens") and prompt_details.text_tokens is not None - else None - ), - audio_tokens=( - prompt_details.audio_tokens - if hasattr(prompt_details, "audio_tokens") and prompt_details.audio_tokens is not None - else None - ), - cache_write_tokens=( - getattr(prompt_details, "cache_write_tokens", None) - or getattr(prompt_details, "cache_creation_tokens", None) - ), + cache_write_tokens: Final = getattr(prompt_details, "cache_write_tokens", None) or getattr( + prompt_details, "cache_creation_tokens", None + ) + input_tokens_details: Final = InputTokensDetails( + cached_tokens=prompt_details.cached_tokens if prompt_details.cached_tokens is not None else 0, + text_tokens=prompt_details.text_tokens, + audio_tokens=prompt_details.audio_tokens, cached_tokens_details=( cached_tokens_details if isinstance(cached_tokens_details, CachedTokensDetails) else None ), ) + if cache_write_tokens is not None: + setattr(input_tokens_details, "cache_write_tokens", cache_write_tokens) + response_usage.input_tokens_details = input_tokens_details # Translate completion_tokens_details to output_tokens_details if hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details is not None: diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 46249e50572..be96c2a4bf5 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -2605,6 +2605,7 @@ class TestUsageTransformation: assert response_usage.input_tokens_details is not None assert response_usage.input_tokens_details.cached_tokens == 5 assert response_usage.input_tokens_details.text_tokens == 8 + assert "cache_write_tokens" not in response_usage.input_tokens_details.model_dump() def test_transform_usage_with_cached_tokens_gemini(self): """Test that cached_tokens from Gemini are properly transformed to input_tokens_details""" @@ -2667,6 +2668,7 @@ class TestUsageTransformation: assert response_usage.input_tokens_details is not None assert response_usage.input_tokens_details.cached_tokens == 100 assert getattr(response_usage.input_tokens_details, "cache_write_tokens", None) == 800 + assert response_usage.input_tokens_details.model_dump()["cache_write_tokens"] == 800 def test_transform_usage_with_reasoning_tokens_gemini(self): """Test that reasoning_tokens from Gemini are properly transformed to output_tokens_details""" From 5737cab258b405689a55e1e6dcaec385673fa572 Mon Sep 17 00:00:00 2001 From: shivam Date: Thu, 10 Sep 2026 22:27:58 +0000 Subject: [PATCH 005/112] fix(cost): cap nested cached modality counts at cached_tokens Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/cost_calculator.py | 2 -- .../litellm_core_utils/llm_cost_calc/utils.py | 16 +++++++++++---- .../llm_cost_calc/test_llm_cost_calc_utils.py | 20 +++++++++++++++++++ 3 files changed, 32 insertions(+), 6 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 8daa2de416b..b865318f3af 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2332,8 +2332,6 @@ def _combine_prompt_tokens_details(combined: Usage, usage: Usage) -> None: if not hasattr(combined, "prompt_tokens_details") or not combined.prompt_tokens_details: combined.prompt_tokens_details = PromptTokensDetailsWrapper() - # Check what keys exist in the model's prompt_tokens_details - # Access model_fields on the class, not the instance, to avoid Pydantic 2.11+ deprecation warnings for attr in _summable_prompt_token_fields(usage.prompt_tokens_details): if ( hasattr(usage.prompt_tokens_details, attr) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 18ef99597a0..dc689ca9618 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -806,9 +806,17 @@ def parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: or None ) cached_tokens_details: Final = getattr(usage.prompt_tokens_details, "cached_tokens_details", None) - cached_text_tokens: Final = _get_token_detail_value(cached_tokens_details, "text_tokens") or 0 - cached_audio_tokens: Final = _get_token_detail_value(cached_tokens_details, "audio_tokens") or 0 - cached_image_tokens: Final = _get_token_detail_value(cached_tokens_details, "image_tokens") or 0 + cached_audio_tokens: Final = min( + _get_token_detail_value(cached_tokens_details, "audio_tokens") or 0, cache_hit_tokens + ) + cached_text_tokens: Final = min( + _get_token_detail_value(cached_tokens_details, "text_tokens") or 0, + cache_hit_tokens - cached_audio_tokens, + ) + cached_image_tokens: Final = min( + _get_token_detail_value(cached_tokens_details, "image_tokens") or 0, + cache_hit_tokens - cached_audio_tokens - cached_text_tokens, + ) text_tokens: Final = max( ( cast(int | None, getattr(usage.prompt_tokens_details, "text_tokens", None)) @@ -852,7 +860,7 @@ def parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: return PromptTokensDetailsResult( cache_hit_tokens=cache_hit_tokens, - cache_hit_audio_tokens=min(cached_audio_tokens, cache_hit_tokens), + cache_hit_audio_tokens=cached_audio_tokens, cache_creation_tokens=cache_creation_tokens, cache_creation_token_details=cache_creation_token_details, text_tokens=text_tokens, diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 3709b526c3b..33825c8dd01 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -5215,3 +5215,23 @@ def test_cached_audio_tokens_fall_back_to_cache_read_input_token_cost() -> None: ) expected = 52 * 4e-6 + 64 * 5e-7 + 39 * 32e-6 + 128 * 5e-7 assert prompt_cost == pytest.approx(expected) + + +def test_cached_audio_tokens_capped_at_cached_tokens(_local_model_cost_map: None) -> None: + """Nested cached_tokens_details exceeding cached_tokens must not over-subtract the audio bucket.""" + usage = Usage( + prompt_tokens=283, + completion_tokens=0, + total_tokens=283, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=116, + audio_tokens=167, + cached_tokens=100, + cached_tokens_details={"audio_tokens": 128}, + ), + ) + + prompt_cost, _ = generic_cost_per_token( + model="gpt-realtime-2", usage=usage, custom_llm_provider="openai" + ) + assert prompt_cost == pytest.approx(116 * 4e-6 + (167 - 100) * 32e-6 + 100 * 4e-7) From a067557dae5c0bb52fdb11176437a39c9a0d9ac3 Mon Sep 17 00:00:00 2001 From: yucheng Date: Fri, 11 Sep 2026 23:16:53 +0000 Subject: [PATCH 006/112] fix(otel): index the opener and the latest prompt turns, not the oldest A value length limit clips the input.value blob, so the per-index keys are the only untruncated copy of a message. Indexing the leading prompt messages dropped the live user turn from every span attribute on long conversations. Keep message 0 and the most recent turns under the same span-wide budget, original indices preserved, reply reservation unchanged Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../otel/mappers/openinference.py | 31 +++++++++++++------ .../integrations/otel/test_otel_v2_emitter.py | 29 ++++++++++++++++- 2 files changed, 50 insertions(+), 10 deletions(-) diff --git a/litellm/integrations/otel/mappers/openinference.py b/litellm/integrations/otel/mappers/openinference.py index 1e2dbf6974d..1fa19b8d4a5 100644 --- a/litellm/integrations/otel/mappers/openinference.py +++ b/litellm/integrations/otel/mappers/openinference.py @@ -92,8 +92,13 @@ class OpenInferenceMapper: return { **collect(self._LLM_CALL_ATTRS, data), **collect(self._BLOB_ATTRS, data), - **self._messages("llm.input_messages", "input.value", data.messages_in, indexed_in), - **self._messages("llm.output_messages", "output.value", outputs, indexed_out), + **self._messages( + "llm.input_messages", + "input.value", + data.messages_in, + self._prompt_positions(len(data.messages_in), indexed_in), + ), + **self._messages("llm.output_messages", "output.value", outputs, range(indexed_out)), **self._tools(data), } @@ -109,18 +114,26 @@ class OpenInferenceMapper: return _MAX_INDEXED_MESSAGES - indexed_out, indexed_out @staticmethod - def _messages(prefix: str, value_key: str, messages: Sequence[object], indexed: int) -> AttributeMap: - """``{prefix}.{idx}.message.*`` keys for the leading ``indexed`` messages + the ``value_key`` blob of all.""" + def _prompt_positions(total: int, indexed: int) -> tuple[int, ...]: + """Which prompt messages get per-index attributes: message 0 and the most recent turns. + + A value length limit clips the ``input.value`` blob, so the system prompt and the + live turn each keep a short key of their own. The middle of a long prompt does not. + """ + if total <= indexed: + return tuple(range(total)) + return (0, *range(total - indexed + 1, total)) + + @staticmethod + def _messages(prefix: str, value_key: str, messages: Sequence[object], positions: Sequence[int]) -> AttributeMap: + """``{prefix}.{idx}.message.*`` keys for the messages at ``positions`` + the ``value_key`` blob of all.""" parsed: Final = [(m.get("role") if isinstance(m, dict) else None, message_content(m)) for m in messages] attrs: Final = drop_none( { key: value - for idx, (role, content) in enumerate(parsed[:indexed]) + for idx, (role, content) in ((idx, parsed[idx]) for idx in positions) for key, value in ( - ( - f"{prefix}.{idx}.message.role", - role if isinstance(role, str) else None, - ), + (f"{prefix}.{idx}.message.role", role if isinstance(role, str) else None), (f"{prefix}.{idx}.message.content", content), ) } diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py b/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py index f571ab7004b..6491ab0f79f 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py @@ -505,7 +505,8 @@ def test_long_conversation_does_not_evict_core_attributes(turns): assert a["llm.input_messages.0.message.content"] == "turn 0" assert a["llm.output_messages.0.message.content"] == "reply 0" - assert f"llm.input_messages.{turns - 1}.message.role" not in a + assert a[f"llm.input_messages.{turns - 1}.message.content"] == f"turn {turns - 1}" + assert f"llm.input_messages.{turns // 2}.message.role" not in a assert len(json.loads(a["input.value"])) == turns assert len(json.loads(a["output.value"])) == 1 assert len(json.loads(a[GenAI.INPUT_MESSAGES])) == turns @@ -520,6 +521,31 @@ def test_short_conversation_keeps_every_message_indexed(): assert a[f"llm.output_messages.{idx}.message.content"] == f"reply {idx}" +def test_indexed_prompt_keeps_opener_and_latest_turns_under_a_value_length_limit(monkeypatch): + """The per-index keys are the only untruncated copy once the SDK clips string values. + + Operators bound attribute sizes with ``OTEL_SPAN_ATTRIBUTE_VALUE_LENGTH_LIMIT``, + which cuts the ``input.value`` blob short. The system prompt and the live turn + then have to survive as their own short keys, whatever the conversation length. + """ + monkeypatch.setenv("OTEL_SPAN_ATTRIBUTE_VALUE_LENGTH_LIMIT", "256") + payload = _conversation_payload(60) + payload["messages"][0] = {"role": "system", "content": "be terse"} + payload["messages"][-1] = {"role": "user", "content": "LATEST-TURN"} + a = _conversation_span(["genai", "openinference"], payload).attributes + + assert len(a["input.value"]) == 256 + assert a["llm.input_messages.0.message.role"] == "system" + assert a["llm.input_messages.0.message.content"] == "be terse" + assert a["llm.input_messages.59.message.role"] == "user" + assert a["llm.input_messages.59.message.content"] == "LATEST-TURN" + assert a["llm.output_messages.0.message.content"] == "reply 0" + assert [int(key.split(".")[2]) for key in a if key.endswith("message.content") and key.startswith("llm.input_")] == [ + 0, + *range(54, 60), + ] + + def test_message_cap_is_shared_across_input_and_output(): """One span-wide allowance covers both directions, and the response always keeps a share. @@ -588,4 +614,5 @@ def test_fully_populated_span_with_every_vocabulary_stays_within_the_attribute_l assert a[f"{LiteLLM.COST_PREFIX}total"] == 0.002 assert a[LiteLLM.TOOLS_DECLARED] == 127 assert a["llm.input_messages.0.message.content"] == "turn 0" + assert a["llm.input_messages.199.message.content"] == "turn 199" assert a["llm.output_messages.0.message.content"] == "reply 0" From a15309dfe820836a41e914228359d7b5becc3744 Mon Sep 17 00:00:00 2001 From: yucheng Date: Fri, 11 Sep 2026 23:32:44 +0000 Subject: [PATCH 007/112] refactor(otel): trim the message cap docstrings to one line each Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../otel/mappers/openinference.py | 13 ++------- litellm/integrations/otel/mappers/utils.py | 13 ++------- .../integrations/otel/test_otel_v2_emitter.py | 29 +++---------------- 3 files changed, 9 insertions(+), 46 deletions(-) diff --git a/litellm/integrations/otel/mappers/openinference.py b/litellm/integrations/otel/mappers/openinference.py index 1fa19b8d4a5..a7e0f1af3ac 100644 --- a/litellm/integrations/otel/mappers/openinference.py +++ b/litellm/integrations/otel/mappers/openinference.py @@ -104,22 +104,13 @@ class OpenInferenceMapper: @staticmethod def _indexed_split(inputs: int, outputs: int) -> tuple[int, int]: - """How many prompt and response messages get per-index attributes. - - Both directions share one span-wide allowance. The response is reserved at - least half of it, so a long prompt can never push the completion off the - span, and the prompt takes whatever the response leaves unused. - """ + """Prompt and response share one allowance; the response is reserved at least half of it.""" indexed_out: Final = min(outputs, max(_MAX_INDEXED_MESSAGES // 2, _MAX_INDEXED_MESSAGES - inputs)) return _MAX_INDEXED_MESSAGES - indexed_out, indexed_out @staticmethod def _prompt_positions(total: int, indexed: int) -> tuple[int, ...]: - """Which prompt messages get per-index attributes: message 0 and the most recent turns. - - A value length limit clips the ``input.value`` blob, so the system prompt and the - live turn each keep a short key of their own. The middle of a long prompt does not. - """ + """Prompt messages that get per-index attributes: message 0 and the most recent turns.""" if total <= indexed: return tuple(range(total)) return (0, *range(total - indexed + 1, total)) diff --git a/litellm/integrations/otel/mappers/utils.py b/litellm/integrations/otel/mappers/utils.py index d8918491720..c023621d2ef 100644 --- a/litellm/integrations/otel/mappers/utils.py +++ b/litellm/integrations/otel/mappers/utils.py @@ -33,17 +33,10 @@ core telemetry no matter how many vocabularies are configured. MAX_MESSAGE_ATTRS_PER_SPAN: Final = DEFAULT_SPAN_ATTRIBUTE_LIMIT // 8 -"""Span-wide ceiling on attributes spent spelling out chat messages per index. +"""Span-wide ceiling on per-index chat message attributes, prompt and response together. -A conversation is the other unbounded family: two attributes per message, for -the prompt and the response alike, on the same span. Past a few dozen turns the -family alone exceeds the span attribute limit and evicts the core telemetry -written before it. The ceiling covers both directions together, since a budget -handed to each direction separately doubles. An eighth is the largest share -that still fits beside the tool ceiling and the core of every vocabulary at -once, request parameters, cost breakdown and identity included. The complete -conversation still rides the JSON blob attributes; only the per-index -convenience keys are capped. +An eighth is the largest share that still fits beside the tool ceiling and the core +of every vocabulary at once. The complete conversation still rides the JSON blobs. """ diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py b/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py index 6491ab0f79f..16fbb242ebd 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py @@ -485,13 +485,7 @@ def _indexed_message_count(attributes, prefix): @pytest.mark.parametrize("turns", [60, 200]) def test_long_conversation_does_not_evict_core_attributes(turns): - """Per-message OpenInference attributes must never crowd core telemetry off the span. - - With content capture on, the OpenInference vocabulary spells every prompt and - response message out as two per-index attributes. A few dozen turns overruns - the OTel SDK's 128-attribute span limit, which evicts oldest-first, so the - ``gen_ai.*`` set written before it is what disappears. - """ + """Per-message OpenInference attributes must never crowd core telemetry off the span.""" span = _conversation_span(["genai", "openinference"], _conversation_payload(turns)) a = span.attributes @@ -522,12 +516,7 @@ def test_short_conversation_keeps_every_message_indexed(): def test_indexed_prompt_keeps_opener_and_latest_turns_under_a_value_length_limit(monkeypatch): - """The per-index keys are the only untruncated copy once the SDK clips string values. - - Operators bound attribute sizes with ``OTEL_SPAN_ATTRIBUTE_VALUE_LENGTH_LIMIT``, - which cuts the ``input.value`` blob short. The system prompt and the live turn - then have to survive as their own short keys, whatever the conversation length. - """ + """The system prompt and the live turn keep their own keys once the SDK clips ``input.value``.""" monkeypatch.setenv("OTEL_SPAN_ATTRIBUTE_VALUE_LENGTH_LIMIT", "256") payload = _conversation_payload(60) payload["messages"][0] = {"role": "system", "content": "be terse"} @@ -547,11 +536,7 @@ def test_indexed_prompt_keeps_opener_and_latest_turns_under_a_value_length_limit def test_message_cap_is_shared_across_input_and_output(): - """One span-wide allowance covers both directions, and the response always keeps a share. - - A long prompt takes what a single reply leaves over, and a many-choice reply - cannot take the whole allowance away from the prompt either. - """ + """One span-wide allowance covers both directions, and the response always keeps a share.""" long_prompt = _conversation_span(["genai", "openinference"], _conversation_payload(60, choices=1)).attributes many_choices = _conversation_span(["genai", "openinference"], _conversation_payload(60, choices=20)).attributes @@ -569,13 +554,7 @@ def test_message_cap_is_shared_across_input_and_output(): def test_fully_populated_span_with_every_vocabulary_stays_within_the_attribute_limit(): - """Every capped family maxed at once still leaves the whole core intact. - - Every vocabulary in the registry plus ``legacy``, every request parameter, - every cost component, a hundred-plus tools, a two-hundred-turn prompt and - twenty choices is the worst case the two span-wide ceilings have to absorb - together. - """ + """Every capped family maxed at once still leaves the whole core intact.""" payload = _conversation_payload( 200, choices=20, From 107b4ec4db64985de0b3651f401b290ea09e81ed Mon Sep 17 00:00:00 2001 From: yassin Date: Sat, 12 Sep 2026 01:12:39 +0000 Subject: [PATCH 008/112] fix(redis): log a timeout streak once per interval instead of one line per cache call Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/caching/redis_cache.py | 140 ++++++++++++------ litellm/constants.py | 3 + tests/test_litellm/caching/test_dual_cache.py | 57 +++++++ .../test_litellm/caching/test_redis_cache.py | 57 +++++++ 4 files changed, 211 insertions(+), 46 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 2c36995c4f8..eaac7ef7b0b 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -15,6 +15,7 @@ import hashlib import inspect import json import logging +import threading import time from collections.abc import Awaitable, Callable, Sequence from contextvars import ContextVar @@ -32,6 +33,7 @@ from litellm.constants import ( REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD, REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT, REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION, + REDIS_TIMEOUT_LOG_INTERVAL, ) from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs from litellm.litellm_core_utils.coroutine_checker import coroutine_checker @@ -404,13 +406,58 @@ class RedisCircuitBreakerOpenError(Exception): pass +class _RedisTimeoutLogThrottle: + """Admits one Redis timeout log line per interval and counts the timeouts it suppressed in between.""" + + def __init__(self, interval: float, clock: Callable[[], float] = time.time) -> None: + self.interval = interval + self._clock = clock + self._lock = threading.Lock() + self._last_logged_at: float | None = None + self._suppressed = 0 + + def admit(self) -> int | None: + """Return the number of timeouts suppressed since the last admitted line, or None to suppress this one.""" + with self._lock: + now: Final = self._clock() + if self._last_logged_at is not None and now - self._last_logged_at < self.interval: + self._suppressed += 1 + return None + suppressed: Final = self._suppressed + self._suppressed = 0 + self._last_logged_at = now + return suppressed + + +_redis_timeout_log_throttle: Final = _RedisTimeoutLogThrottle(REDIS_TIMEOUT_LOG_INTERVAL) + + def log_redis_failure( logger: logging.Logger, level: int, message: str, exc: BaseException, with_traceback: bool = False ) -> None: if isinstance(exc, RedisCircuitBreakerOpenError): - logger.debug("%s: %s", message, exc) + logger.debug("%s: %s", message, exc, stacklevel=2) return - logger.log(level, "%s: %s", message, exc, exc_info=exc if with_traceback else None) + exc_info: Final = exc if with_traceback else None + if not _is_redis_timeout_failure(exc): + logger.log(level, "%s: %s", message, exc, exc_info=exc_info, stacklevel=2) + return + suppressed: Final = _redis_timeout_log_throttle.admit() + if suppressed is None: + logger.debug("%s: %s", message, exc, stacklevel=2) + return + if suppressed == 0: + logger.log(level, "%s: %s", message, exc, exc_info=exc_info, stacklevel=2) + return + logger.log( + level, + "%s: %s (%d more Redis timeouts since the previous Redis timeout line were logged at DEBUG)", + message, + exc, + suppressed, + exc_info=exc_info, + stacklevel=2, + ) @dataclass(frozen=True, slots=True) @@ -783,10 +830,8 @@ class RedisCache(BaseCache): ## LOGGING ## end_time = time.time() _duration = end_time - start_time - verbose_logger.error( - "LiteLLM Redis Caching: increment_cache() - Got exception from REDIS %s, Writing value=%s", - str(e), - value, + log_redis_failure( + verbose_logger, logging.ERROR, "LiteLLM Redis Caching: increment_cache() - Got exception from REDIS", e ) raise e @@ -992,11 +1037,8 @@ class RedisCache(BaseCache): call_type=f"async_set_cache <- {_get_call_stack_info()}", ) ) - verbose_logger.error( - "LiteLLM Redis Caching: async set() - Got exception from REDIS %s, key=%r, value=%r", - str(e), - key, - value, + log_redis_failure( + verbose_logger, logging.ERROR, "LiteLLM Redis Caching: async set() - Got exception from REDIS", e ) raise e @@ -1044,10 +1086,8 @@ class RedisCache(BaseCache): event_metadata={"key": key}, ) ) - verbose_logger.error( - "LiteLLM Redis Caching: async set() - Got exception from REDIS %s, Writing value=%s", - str(e), - value, + log_redis_failure( + verbose_logger, logging.ERROR, "LiteLLM Redis Caching: async set() - Got exception from REDIS", e ) _record_swallowed_redis_failure(self._circuit_breaker, e) @@ -1094,7 +1134,6 @@ class RedisCache(BaseCache): start_time: Final = time.time() print_verbose(f"Set Async Redis Cache: key list: {cache_list}\nttl={ttl}, redis_version={self.redis_version}") - cache_value: Final = None try: async with _redis_client.pipeline(transaction=False) as pipe: results: Final = await self._pipeline_helper(pipe, cache_list, ttl) @@ -1131,10 +1170,11 @@ class RedisCache(BaseCache): ) ) - verbose_logger.error( - "LiteLLM Redis Caching: async set_cache_pipeline() - Got exception from REDIS %s, Writing value=%s", - str(e), - cache_value, + log_redis_failure( + verbose_logger, + logging.ERROR, + "LiteLLM Redis Caching: async set_cache_pipeline() - Got exception from REDIS", + e, ) _record_swallowed_redis_failure(self._circuit_breaker, e) @@ -1177,10 +1217,8 @@ class RedisCache(BaseCache): ) ) # NON blocking - notify users Redis is throwing an exception - verbose_logger.error( - "LiteLLM Redis Caching: async set() - Got exception from REDIS %s, Writing value=%s", - str(e), - value, + log_redis_failure( + verbose_logger, logging.ERROR, "LiteLLM Redis Caching: async set() - Got exception from REDIS", e ) raise e @@ -1216,10 +1254,11 @@ class RedisCache(BaseCache): ) ) # NON blocking - notify users Redis is throwing an exception - verbose_logger.error( - "LiteLLM Redis Caching: async set_cache_sadd() - Got exception from REDIS %s, Writing value=%s", - str(e), - value, + log_redis_failure( + verbose_logger, + logging.ERROR, + "LiteLLM Redis Caching: async set_cache_sadd() - Got exception from REDIS", + e, ) _record_swallowed_redis_failure(self._circuit_breaker, e) @@ -1288,10 +1327,11 @@ class RedisCache(BaseCache): parent_otel_span=parent_otel_span, ) ) - verbose_logger.error( - "LiteLLM Redis Caching: async async_increment() - Got exception from REDIS %s, Writing value=%s", - str(e), - value, + log_redis_failure( + verbose_logger, + logging.ERROR, + "LiteLLM Redis Caching: async async_increment() - Got exception from REDIS", + e, ) raise e @@ -1377,7 +1417,9 @@ class RedisCache(BaseCache): print_verbose(f"Got Redis Cache: key: {key}, cached_response {cached_response}") return self._get_cache_logic(cached_response=cached_response) except Exception as e: - verbose_logger.error("litellm.caching.caching: get() - Got exception from REDIS: %s", e) + log_redis_failure( + verbose_logger, logging.ERROR, "litellm.caching.caching: get() - Got exception from REDIS", e + ) _record_swallowed_redis_failure(self._circuit_breaker, e) def _run_redis_mget_operation(self, keys: list[str]) -> Sequence[bytes | str | None]: @@ -1455,7 +1497,7 @@ class RedisCache(BaseCache): end_time=failed_at, parent_otel_span=parent_otel_span, ) - verbose_logger.error("Error occurred in batch get cache - %s", e) + log_redis_failure(verbose_logger, logging.ERROR, "Error occurred in batch get cache", e) _record_swallowed_redis_failure(self._circuit_breaker, e) return key_value_dict @@ -1574,7 +1616,7 @@ class RedisCache(BaseCache): parent_otel_span=parent_otel_span, ) ) - verbose_logger.error("Error occurred in async batch get cache - %s", e) + log_redis_failure(verbose_logger, logging.ERROR, "Error occurred in async batch get cache", e) _record_swallowed_redis_failure(self._circuit_breaker, e) return key_value_dict @@ -1799,9 +1841,11 @@ class RedisCache(BaseCache): parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs), ) ) - verbose_logger.error( - "LiteLLM Redis Caching: async increment_pipeline() - Got exception from REDIS %s", - str(e), + log_redis_failure( + verbose_logger, + logging.ERROR, + "LiteLLM Redis Caching: async increment_pipeline() - Got exception from REDIS", + e, ) raise e @@ -1878,7 +1922,7 @@ class RedisCache(BaseCache): call_type=f"async_rpush <- {_get_call_stack_info()}", ) ) - verbose_logger.error("LiteLLM Redis Cache RPUSH: - Got exception from REDIS : %s", e) + log_redis_failure(verbose_logger, logging.ERROR, "LiteLLM Redis Cache RPUSH: - Got exception from REDIS", e) raise e async def _pipeline_rpush_helper( @@ -1946,9 +1990,11 @@ class RedisCache(BaseCache): call_type=f"async_rpush_pipeline <- {_get_call_stack_info()}", ) ) - verbose_logger.error( - "LiteLLM Redis Caching: async_rpush_pipeline() - Got exception from REDIS %s", - str(e), + log_redis_failure( + verbose_logger, + logging.ERROR, + "LiteLLM Redis Caching: async_rpush_pipeline() - Got exception from REDIS", + e, ) raise e @@ -2024,7 +2070,7 @@ class RedisCache(BaseCache): call_type=f"async_lpop <- {_get_call_stack_info()}", ) ) - verbose_logger.error("LiteLLM Redis Cache LPOP: - Got exception from REDIS : %s", e) + log_redis_failure(verbose_logger, logging.ERROR, "LiteLLM Redis Cache LPOP: - Got exception from REDIS", e) raise e async def _pipeline_lpop_helper( @@ -2135,8 +2181,10 @@ class RedisCache(BaseCache): call_type=f"async_lpop_pipeline <- {_get_call_stack_info()}", ) ) - verbose_logger.error( - "LiteLLM Redis Caching: async_lpop_pipeline() - Got exception from REDIS %s", - str(e), + log_redis_failure( + verbose_logger, + logging.ERROR, + "LiteLLM Redis Caching: async_lpop_pipeline() - Got exception from REDIS", + e, ) raise e diff --git a/litellm/constants.py b/litellm/constants.py index 6b984c2673c..a32551b4480 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -459,6 +459,9 @@ REDIS_CIRCUIT_BREAKER_ENABLED: Final = os.getenv("REDIS_CIRCUIT_BREAKER_ENABLED" # minimum seconds a timeout-only failure streak must span before it can open the breaker, # so one event-loop stall timing out many queued calls at once does not trip it REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION: Final = float(os.getenv("REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION", 5.0)) +# seconds between Redis timeout log lines: the first timeout of a streak logs at the caller's level, +# later ones log at DEBUG until the interval passes and one line summarizes how many were suppressed +REDIS_TIMEOUT_LOG_INTERVAL: Final = float(os.getenv("REDIS_TIMEOUT_LOG_INTERVAL", "5.0")) # Seconds of idle before a Redis cluster connection is validated with a PING and # reconnected if dead, so a connection silently dropped by a cluster restart # (e.g. ElastiCache Serverless maintenance) is not reused while broken diff --git a/tests/test_litellm/caching/test_dual_cache.py b/tests/test_litellm/caching/test_dual_cache.py index 4c9068722b8..850fa14106b 100644 --- a/tests/test_litellm/caching/test_dual_cache.py +++ b/tests/test_litellm/caching/test_dual_cache.py @@ -704,3 +704,60 @@ async def test_open_breaker_keeps_async_batch_read_memory_hits_and_releases_rese assert list(await cache.async_batch_get_cache(["k1", "k2"])) == ["v1", None] assert "k2" not in cache.last_redis_batch_access_time + + +@pytest.mark.asyncio +async def test_redis_timeouts_falling_back_to_memory_log_once_per_interval(caplog, monkeypatch): + """The in-memory fallback WARNING must not repeat for every timed-out increment during a blip. + + The rate limiter's pipeline increments and the dual cache increments each logged a WARNING per + call while Redis timed out, hundreds of lines per second before the breaker opened. The first + timeout of a streak keeps its WARNING, the rest are DEBUG until the summary interval passes. + """ + from redis.exceptions import TimeoutError as RedisTimeoutError + + from litellm.caching import redis_cache as redis_cache_module + from litellm.caching.redis_cache import _RedisTimeoutLogThrottle + + clock = MagicMock(return_value=1_000.0) + monkeypatch.setattr( + redis_cache_module, "_redis_timeout_log_throttle", _RedisTimeoutLogThrottle(interval=5.0, clock=clock) + ) + + class _TimingOutRedis: + async def async_increment_pipeline(self, increment_list, **kwargs): + raise RedisTimeoutError("Timeout reading from 127.0.0.1:6379") + + async def async_increment(self, key, value, **kwargs): + raise RedisTimeoutError("Timeout reading from 127.0.0.1:6379") + + cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=_TimingOutRedis()) # pyright: ignore[reportArgumentType] # duck-typed Redis double + increments = [RedisPipelineIncrementOperation(key="k", increment_value=1.0, ttl=60)] + + with caplog.at_level(logging.DEBUG, logger="LiteLLM"): + for _ in range(100): + await cache.async_increment_cache_pipeline(increment_list=increments) + await cache.async_increment_cache("k", 1.0) + + visible = [r for r in caplog.records if r.levelno >= logging.WARNING] + assert [(r.levelno, r.getMessage()) for r in visible] == [ + ( + logging.WARNING, + "Redis async_increment_cache_pipeline failed, falling back to in-memory result:" + " Timeout reading from 127.0.0.1:6379", + ) + ] + assert visible[0].filename == "dual_cache.py" + assert sum("Timeout reading from" in r.getMessage() for r in caplog.records) == 200 + + caplog.clear() + clock.return_value += 5.0 + with caplog.at_level(logging.DEBUG, logger="LiteLLM"): + await cache.async_increment_cache("k", 1.0) + assert [(r.levelno, r.getMessage()) for r in caplog.records] == [ + ( + logging.WARNING, + "Redis async_increment_cache failed, falling back to in-memory result: Timeout reading from 127.0.0.1:6379" + " (199 more Redis timeouts since the previous Redis timeout line were logged at DEBUG)", + ) + ] diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index bcae33b976e..d0974b2420c 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -1202,3 +1202,60 @@ async def test_a_probe_overtaken_by_a_later_outage_leaves_the_breaker_to_the_new new_probe_release.set() assert await new_probe == "new probe" assert breaker._state == breaker.CLOSED + + +def test_timeouts_during_a_blip_log_once_per_interval_not_once_per_call(sync_batch_redis_cache, caplog, monkeypatch): + """A Redis latency blip must not write one ERROR line per timed-out cache call. + + Before the breaker opens (up to REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION of timeouts) every + cache operation logged its own ERROR or WARNING line, so one single-worker proxy wrote + ~1100 lines in 5 s at LITELLM_LOG=WARNING. A timeout streak now logs its first failure, then + one summary line per REDIS_TIMEOUT_LOG_INTERVAL carrying the count of suppressed timeouts, + while every timeout stays visible at DEBUG. Hard connectivity failures keep their per-call line. + """ + import logging + + from redis.exceptions import TimeoutError as RedisTimeoutError + + from litellm.caching import redis_cache as redis_cache_module + from litellm.caching.redis_cache import _RedisTimeoutLogThrottle + + clock = MagicMock(return_value=1_000.0) + monkeypatch.setattr( + redis_cache_module, "_redis_timeout_log_throttle", _RedisTimeoutLogThrottle(interval=5.0, clock=clock) + ) + sync_batch_redis_cache.redis_client.get.side_effect = RedisTimeoutError("Timeout reading from 127.0.0.1:6379") + sync_batch_redis_cache.redis_client.mget.side_effect = RedisTimeoutError("Timeout reading from 127.0.0.1:6379") + + with caplog.at_level(logging.DEBUG, logger="LiteLLM"): + for _ in range(200): + assert sync_batch_redis_cache.get_cache("lit7520") is None + assert sync_batch_redis_cache.batch_get_cache(key_list=["lit7520"]) == {} + + timeout_records = [r for r in caplog.records if "Timeout reading from" in r.getMessage()] + assert len(timeout_records) == 201, "every timeout must stay visible at DEBUG" + assert [r.getMessage() for r in timeout_records if r.levelno >= logging.WARNING] == [ + "litellm.caching.caching: get() - Got exception from REDIS: Timeout reading from 127.0.0.1:6379" + ] + assert timeout_records[0].levelno == logging.ERROR + assert timeout_records[0].filename == "redis_cache.py" + assert timeout_records[0].lineno != timeout_records[-1].lineno, "the record must point at the cache operation" + + caplog.clear() + clock.return_value += 5.0 + with caplog.at_level(logging.DEBUG, logger="LiteLLM"): + assert sync_batch_redis_cache.batch_get_cache(key_list=["lit7520"]) == {} + assert [(r.levelno, r.getMessage()) for r in caplog.records] == [ + ( + logging.ERROR, + "Error occurred in batch get cache: Timeout reading from 127.0.0.1:6379" + " (200 more Redis timeouts since the previous Redis timeout line were logged at DEBUG)", + ) + ] + + caplog.clear() + sync_batch_redis_cache.redis_client.get.side_effect = OSError("redis unavailable") + with caplog.at_level(logging.DEBUG, logger="LiteLLM"): + for _ in range(3): + assert sync_batch_redis_cache.get_cache("lit7520") is None + assert [r.levelno for r in caplog.records if "redis unavailable" in r.getMessage()] == [logging.ERROR] * 3 From 9c84e98fb22bd0f6e2c359f335bbc329181bb8bd Mon Sep 17 00:00:00 2001 From: yassin Date: Sat, 12 Sep 2026 01:25:32 +0000 Subject: [PATCH 009/112] fix(proxy): treat a Redis timeout in spend counter increments as an already-logged cache failure The cost tracking callback logged its own ERROR with a traceback for every request whose spend counter increment timed out, on top of the cache layer's throttled line. Timeouts now take the same path as breaker-open refusals: invalidate the counters and return. Also exposes is_redis_timeout_failure publicly for that caller and drops the comment on the new constant Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/caching/redis_cache.py | 10 +++---- litellm/constants.py | 2 -- litellm/proxy/proxy_server.py | 4 +-- .../test_litellm/caching/test_redis_cache.py | 26 +++++++++---------- .../proxy/proxy_server/test_spend_counters.py | 22 ++++++++++++++++ 5 files changed, 42 insertions(+), 22 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index eaac7ef7b0b..7a3e689a667 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -330,7 +330,7 @@ def _redis_timeout_error_types() -> tuple[type, ...]: return (RedisTimeoutError, TimeoutError) -def _is_redis_timeout_failure(exc: BaseException) -> bool: +def is_redis_timeout_failure(exc: BaseException) -> bool: return isinstance(exc, _redis_timeout_error_types()) @@ -398,7 +398,7 @@ def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseExcep """ if not _is_redis_health_failure(exc): return - breaker.record_failure(is_timeout=_is_redis_timeout_failure(exc)) + breaker.record_failure(is_timeout=is_redis_timeout_failure(exc)) _swallowed_redis_failures.set(_swallowed_redis_failures.get() + 1) @@ -439,7 +439,7 @@ def log_redis_failure( logger.debug("%s: %s", message, exc, stacklevel=2) return exc_info: Final = exc if with_traceback else None - if not _is_redis_timeout_failure(exc): + if not is_redis_timeout_failure(exc): logger.log(level, "%s: %s", message, exc, exc_info=exc_info, stacklevel=2) return suppressed: Final = _redis_timeout_log_throttle.admit() @@ -504,7 +504,7 @@ async def _run_under_circuit_breaker( result: Final = await call() except Exception as e: if _is_redis_health_failure(e): - breaker.record_failure(is_timeout=_is_redis_timeout_failure(e)) + breaker.record_failure(is_timeout=is_redis_timeout_failure(e)) raise _exit_circuit_breaker(breaker, admission) return result @@ -521,7 +521,7 @@ def _run_under_circuit_breaker_sync( result: Final = call() except Exception as e: if _is_redis_health_failure(e): - breaker.record_failure(is_timeout=_is_redis_timeout_failure(e)) + breaker.record_failure(is_timeout=is_redis_timeout_failure(e)) raise _exit_circuit_breaker(breaker, admission) return result diff --git a/litellm/constants.py b/litellm/constants.py index a32551b4480..60e1c682238 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -459,8 +459,6 @@ REDIS_CIRCUIT_BREAKER_ENABLED: Final = os.getenv("REDIS_CIRCUIT_BREAKER_ENABLED" # minimum seconds a timeout-only failure streak must span before it can open the breaker, # so one event-loop stall timing out many queued calls at once does not trip it REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION: Final = float(os.getenv("REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION", 5.0)) -# seconds between Redis timeout log lines: the first timeout of a streak logs at the caller's level, -# later ones log at DEBUG until the interval passes and one line summarizes how many were suppressed REDIS_TIMEOUT_LOG_INTERVAL: Final = float(os.getenv("REDIS_TIMEOUT_LOG_INTERVAL", "5.0")) # Seconds of idle before a Redis cluster connection is validated with a PING and # reconnected if dead, so a connection silently dropped by a cluster restart diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c8fd7edfed6..318ea96dbcb 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -251,7 +251,7 @@ import litellm._redis from litellm import Router from litellm._logging import _redact_string, verbose_proxy_logger, verbose_router_logger from litellm.caching.caching import DualCache, RedisCache -from litellm.caching.redis_cache import RedisCircuitBreakerOpenError +from litellm.caching.redis_cache import RedisCircuitBreakerOpenError, is_redis_timeout_failure from litellm.caching.redis_cluster_cache import RedisClusterCache from litellm.constants import ( _REALTIME_BODY_CACHE_SIZE, @@ -3411,7 +3411,7 @@ async def _apply_spend_counter_increments(pending: Sequence[_PendingSpendIncreme results: Final = await redis_cache.async_increment_pipeline(increment_list=increment_list) except Exception as e: await asyncio.gather(*(_invalidate_spend_counter(counter_key=item.counter_key) for item in pending)) - if isinstance(e, RedisCircuitBreakerOpenError): + if isinstance(e, RedisCircuitBreakerOpenError) or is_redis_timeout_failure(e): return raise for item, current_value in zip(pending, results or ()): diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index d0974b2420c..5840f450ac6 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -977,17 +977,17 @@ async def test_stale_timeout_does_not_let_sub_threshold_hard_failures_open_the_b from redis.exceptions import ConnectionError as RedisConnectionError from redis.exceptions import TimeoutError as RedisTimeoutError - from litellm.caching.redis_cache import RedisCircuitBreaker, _is_redis_timeout_failure + from litellm.caching.redis_cache import RedisCircuitBreaker, is_redis_timeout_failure breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, timeout_min_duration=0.05) - breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisTimeoutError("read timed out"))) + breaker.record_failure(is_timeout=is_redis_timeout_failure(RedisTimeoutError("read timed out"))) await asyncio.sleep(0.06) for _ in range(breaker.failure_threshold - 1): - breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisConnectionError("refused"))) + breaker.record_failure(is_timeout=is_redis_timeout_failure(RedisConnectionError("refused"))) assert breaker.is_open() is False, "2 hard failures and 1 stale timeout are below both thresholds" - breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisConnectionError("refused"))) + breaker.record_failure(is_timeout=is_redis_timeout_failure(RedisConnectionError("refused"))) assert breaker.is_open() is True, "the threshold-th hard failure must still open it" @@ -999,19 +999,19 @@ async def test_hard_failure_resets_timeout_streak_so_a_later_burst_must_earn_its from redis.exceptions import ConnectionError as RedisConnectionError from redis.exceptions import TimeoutError as RedisTimeoutError - from litellm.caching.redis_cache import RedisCircuitBreaker, _is_redis_timeout_failure + from litellm.caching.redis_cache import RedisCircuitBreaker, is_redis_timeout_failure breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, timeout_min_duration=0.05) - breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisTimeoutError("read timed out"))) - breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisConnectionError("refused"))) + breaker.record_failure(is_timeout=is_redis_timeout_failure(RedisTimeoutError("read timed out"))) + breaker.record_failure(is_timeout=is_redis_timeout_failure(RedisConnectionError("refused"))) await asyncio.sleep(0.06) for _ in range(breaker.failure_threshold): - breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisTimeoutError("read timed out"))) + breaker.record_failure(is_timeout=is_redis_timeout_failure(RedisTimeoutError("read timed out"))) assert breaker.is_open() is False, "the burst is instantaneous, so the duration gate must hold it closed" await asyncio.sleep(0.06) - breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisTimeoutError("read timed out"))) + breaker.record_failure(is_timeout=is_redis_timeout_failure(RedisTimeoutError("read timed out"))) assert breaker.is_open() is True, "the same run of timeouts persisting past the duration must open it" @@ -1022,7 +1022,7 @@ async def test_breaker_metrics_track_state_and_failure_class(): from redis.exceptions import ConnectionError as RedisConnectionError from redis.exceptions import TimeoutError as RedisTimeoutError - from litellm.caching.redis_cache import RedisCircuitBreaker, _is_redis_timeout_failure + from litellm.caching.redis_cache import RedisCircuitBreaker, is_redis_timeout_failure def sample(name, labels=None): return REGISTRY.get_sample_value(name, labels) or 0.0 @@ -1034,9 +1034,9 @@ async def test_breaker_metrics_track_state_and_failure_class(): closed_gauge_before = sample("litellm_redis_circuit_breaker_state", {"state": "closed"}) breaker = RedisCircuitBreaker(failure_threshold=2, recovery_timeout=60, timeout_min_duration=5.0) - breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisTimeoutError("t"))) - breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisConnectionError("refused"))) - breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisConnectionError("refused"))) + breaker.record_failure(is_timeout=is_redis_timeout_failure(RedisTimeoutError("t"))) + breaker.record_failure(is_timeout=is_redis_timeout_failure(RedisConnectionError("refused"))) + breaker.record_failure(is_timeout=is_redis_timeout_failure(RedisConnectionError("refused"))) assert sample("litellm_redis_circuit_breaker_failures_total", {"failure_class": "timeout"}) == timeout_before + 1 assert sample("litellm_redis_circuit_breaker_failures_total", {"failure_class": "connectivity"}) == hard_before + 2 diff --git a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py index 2f47736a398..19b5a11af33 100644 --- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py +++ b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py @@ -1180,6 +1180,28 @@ async def test_apply_spend_counter_increments_open_breaker_invalidates_and_retur fake_cache.in_memory_cache.set_cache.assert_not_called() +@pytest.mark.asyncio +async def test_apply_spend_counter_increments_redis_timeout_invalidates_and_returns(monkeypatch): + """A Redis timeout is the streak the breaker is already counting and the cache layer already logged. + + Re-raising it sent every request in the pre-open window through the cost callback's error + path, which logged a traceback and fired the failed-tracking alert once per request. + """ + from redis.exceptions import TimeoutError as RedisTimeoutError + + fake_cache = _make_spend_counter_cache() + fake_cache.redis_cache.async_increment_pipeline = AsyncMock( + side_effect=RedisTimeoutError("Timeout reading from 127.0.0.1:6379") + ) + monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + + await ps._apply_spend_counter_increments(_two_pending_increments()) + + deleted_keys = sorted(call.kwargs["key"] for call in fake_cache.in_memory_cache.delete_cache.call_args_list) + assert deleted_keys == ["spend:key:k", "spend:team:t"] + fake_cache.in_memory_cache.set_cache.assert_not_called() + + @pytest.mark.asyncio async def test_apply_spend_counter_increments_other_redis_error_invalidates_and_raises(monkeypatch): fake_cache = _make_spend_counter_cache() From f681a978f06baa13da0f0c24f7b1ac3a20d9a02a Mon Sep 17 00:00:00 2001 From: yassin Date: Sat, 12 Sep 2026 01:37:36 +0000 Subject: [PATCH 010/112] fix(redis): use a monotonic clock for the timeout log throttle and trim test docstrings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/caching/redis_cache.py | 2 +- tests/test_litellm/caching/test_dual_cache.py | 7 +------ tests/test_litellm/caching/test_redis_cache.py | 9 +-------- .../proxy/proxy_server/test_spend_counters.py | 6 +----- 4 files changed, 4 insertions(+), 20 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 7a3e689a667..e5e27d1e02b 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -409,7 +409,7 @@ class RedisCircuitBreakerOpenError(Exception): class _RedisTimeoutLogThrottle: """Admits one Redis timeout log line per interval and counts the timeouts it suppressed in between.""" - def __init__(self, interval: float, clock: Callable[[], float] = time.time) -> None: + def __init__(self, interval: float, clock: Callable[[], float] = time.monotonic) -> None: self.interval = interval self._clock = clock self._lock = threading.Lock() diff --git a/tests/test_litellm/caching/test_dual_cache.py b/tests/test_litellm/caching/test_dual_cache.py index 850fa14106b..6f29be00b30 100644 --- a/tests/test_litellm/caching/test_dual_cache.py +++ b/tests/test_litellm/caching/test_dual_cache.py @@ -708,12 +708,7 @@ async def test_open_breaker_keeps_async_batch_read_memory_hits_and_releases_rese @pytest.mark.asyncio async def test_redis_timeouts_falling_back_to_memory_log_once_per_interval(caplog, monkeypatch): - """The in-memory fallback WARNING must not repeat for every timed-out increment during a blip. - - The rate limiter's pipeline increments and the dual cache increments each logged a WARNING per - call while Redis timed out, hundreds of lines per second before the breaker opened. The first - timeout of a streak keeps its WARNING, the rest are DEBUG until the summary interval passes. - """ + """The first fallback WARNING of a timeout streak logs, the rest stay at DEBUG until the summary.""" from redis.exceptions import TimeoutError as RedisTimeoutError from litellm.caching import redis_cache as redis_cache_module diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index 5840f450ac6..4ca33894aed 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -1205,14 +1205,7 @@ async def test_a_probe_overtaken_by_a_later_outage_leaves_the_breaker_to_the_new def test_timeouts_during_a_blip_log_once_per_interval_not_once_per_call(sync_batch_redis_cache, caplog, monkeypatch): - """A Redis latency blip must not write one ERROR line per timed-out cache call. - - Before the breaker opens (up to REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION of timeouts) every - cache operation logged its own ERROR or WARNING line, so one single-worker proxy wrote - ~1100 lines in 5 s at LITELLM_LOG=WARNING. A timeout streak now logs its first failure, then - one summary line per REDIS_TIMEOUT_LOG_INTERVAL carrying the count of suppressed timeouts, - while every timeout stays visible at DEBUG. Hard connectivity failures keep their per-call line. - """ + """A timeout streak logs its first failure plus one summary per interval; other failures log per call.""" import logging from redis.exceptions import TimeoutError as RedisTimeoutError diff --git a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py index 19b5a11af33..4a1fc389d3e 100644 --- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py +++ b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py @@ -1182,11 +1182,7 @@ async def test_apply_spend_counter_increments_open_breaker_invalidates_and_retur @pytest.mark.asyncio async def test_apply_spend_counter_increments_redis_timeout_invalidates_and_returns(monkeypatch): - """A Redis timeout is the streak the breaker is already counting and the cache layer already logged. - - Re-raising it sent every request in the pre-open window through the cost callback's error - path, which logged a traceback and fired the failed-tracking alert once per request. - """ + """A Redis timeout invalidates the counters and returns without reaching the cost callback's error path.""" from redis.exceptions import TimeoutError as RedisTimeoutError fake_cache = _make_spend_counter_cache() From 28f2d1f0168aa31639a23447d391516129267069 Mon Sep 17 00:00:00 2001 From: yassin Date: Sat, 12 Sep 2026 01:58:36 +0000 Subject: [PATCH 011/112] test(redis): cover the write and list timeout paths going through the shared log throttle Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_litellm/caching/test_redis_cache.py | 57 +++++++++++++++++++ 1 file changed, 57 insertions(+) diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index 4ca33894aed..5ea21dae539 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -1252,3 +1252,60 @@ def test_timeouts_during_a_blip_log_once_per_interval_not_once_per_call(sync_bat for _ in range(3): assert sync_batch_redis_cache.get_cache("lit7520") is None assert [r.levelno for r in caplog.records if "redis unavailable" in r.getMessage()] == [logging.ERROR] * 3 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "call_method", + [ + pytest.param(lambda c: c.async_set_cache_pipeline([("lit7520", "v")]), id="async_set_cache_pipeline"), + pytest.param(lambda c: c.async_set_cache_sadd("lit7520", ["v"], ttl=None), id="async_set_cache_sadd"), + pytest.param(lambda c: c.async_increment("lit7520", 1.0), id="async_increment"), + pytest.param( + lambda c: c.async_increment_pipeline([{"key": "lit7520", "increment_value": 1.0, "ttl": 60}]), + id="async_increment_pipeline", + ), + pytest.param(lambda c: c.async_rpush("lit7520", ["v"]), id="async_rpush"), + pytest.param( + lambda c: c.async_rpush_pipeline([{"key": "lit7520", "values": ["v"]}]), id="async_rpush_pipeline" + ), + pytest.param(lambda c: c.async_lpop("lit7520"), id="async_lpop"), + pytest.param(lambda c: c.async_lpop_pipeline([{"key": "lit7520", "count": 1}]), id="async_lpop_pipeline"), + ], +) +async def test_write_path_timeouts_inside_the_interval_stay_at_debug(call_method, caplog, monkeypatch, redis_no_ping): + """A write or list operation timing out mid-streak is counted by the throttle instead of logging its own ERROR.""" + import contextlib + import logging + + from redis.exceptions import TimeoutError as RedisTimeoutError + + from litellm.caching import redis_cache as redis_cache_module + from litellm.caching.redis_cache import _RedisTimeoutLogThrottle + + clock = MagicMock(return_value=1_000.0) + throttle = _RedisTimeoutLogThrottle(interval=5.0, clock=clock) + assert throttle.admit() == 0 + monkeypatch.setattr(redis_cache_module, "_redis_timeout_log_throttle", throttle) + + timeout = RedisTimeoutError("Timeout reading from 127.0.0.1:6379") + client = MagicMock() + client.pipeline.return_value.__aenter__.side_effect = timeout + client.sadd = AsyncMock(side_effect=timeout) + client.incrbyfloat = AsyncMock(side_effect=timeout) + client.rpush = AsyncMock(side_effect=timeout) + client.lpop = AsyncMock(side_effect=timeout) + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + cache = RedisCache() + + with ( + patch.object(cache, "init_async_client", return_value=client), + caplog.at_level(logging.DEBUG, logger="LiteLLM"), + ): + with contextlib.suppress(RedisTimeoutError): + await call_method(cache) + + timeout_records = [r for r in caplog.records if "Timeout reading from" in r.getMessage()] + assert [(r.levelno, r.filename) for r in timeout_records] == [(logging.DEBUG, "redis_cache.py")] + clock.return_value += 5.0 + assert throttle.admit() == 1 From 3c2342bfd32edddadfa3429ddefa159fdd9e32c2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 12 Sep 2026 15:05:41 -0700 Subject: [PATCH 012/112] refactor(cost): return a new prompt token details wrapper when combining usage --- litellm/cost_calculator.py | 55 +++++++++++++++++++------------------- 1 file changed, 27 insertions(+), 28 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index b865318f3af..50118af8a30 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2326,33 +2326,29 @@ def _combine_cached_tokens_details( ) -def _combine_prompt_tokens_details(combined: Usage, usage: Usage) -> None: - if not (hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details): - return - if not hasattr(combined, "prompt_tokens_details") or not combined.prompt_tokens_details: - combined.prompt_tokens_details = PromptTokensDetailsWrapper() - - for attr in _summable_prompt_token_fields(usage.prompt_tokens_details): - if ( - hasattr(usage.prompt_tokens_details, attr) - and not attr.startswith("_") - and not callable(_attribute_value(usage.prompt_tokens_details, attr)) - ): - current_val = getattr(combined.prompt_tokens_details, attr, 0) or 0 - new_val = getattr(usage.prompt_tokens_details, attr, 0) or 0 - if new_val is not None and isinstance(new_val, (int, float)): - setattr( - combined.prompt_tokens_details, - attr, - current_val + new_val, - ) - - new_cached_tokens_details: Final = getattr(usage.prompt_tokens_details, "cached_tokens_details", None) - if isinstance(new_cached_tokens_details, CachedTokensDetails): - combined.prompt_tokens_details.cached_tokens_details = _combine_cached_tokens_details( - getattr(combined.prompt_tokens_details, "cached_tokens_details", None), - new_cached_tokens_details, - ) +def _combine_prompt_tokens_details( + current: PromptTokensDetailsWrapper | None, new: PromptTokensDetailsWrapper +) -> PromptTokensDetailsWrapper: + base: Final = current if current is not None else PromptTokensDetailsWrapper() + base_values: Final = MappingProxyType( + {attr: getattr(base, attr) for attr in type(base).model_fields if hasattr(base, attr)} + ) + summed: Final = MappingProxyType( + { + attr: (getattr(base, attr, 0) or 0) + (getattr(new, attr) or 0) + for attr in _summable_prompt_token_fields(new) + if hasattr(new, attr) and isinstance(getattr(new, attr) or 0, (int, float)) + } + ) + new_cached_tokens_details: Final = getattr(new, "cached_tokens_details", None) + cached_tokens_details: Final = ( + _combine_cached_tokens_details(getattr(base, "cached_tokens_details", None), new_cached_tokens_details) + if isinstance(new_cached_tokens_details, CachedTokensDetails) + else getattr(base, "cached_tokens_details", None) + ) + return PromptTokensDetailsWrapper( + **MappingProxyType({**base_values, **summed, "cached_tokens_details": cached_tokens_details}) + ) class BaseTokenUsageProcessor: @@ -2381,7 +2377,10 @@ class BaseTokenUsageProcessor: and isinstance(current_val, (int, float)) ): setattr(combined, attr, current_val + new_val) - _combine_prompt_tokens_details(combined, usage) + if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details: + combined.prompt_tokens_details = _combine_prompt_tokens_details( + getattr(combined, "prompt_tokens_details", None), usage.prompt_tokens_details + ) # Handle nested completion_tokens_details if hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details: From e4b05883627555fd72db16f0ecf376b75a6d93a3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 12 Sep 2026 15:31:53 -0700 Subject: [PATCH 013/112] fix(cost): carry cache_read_input_audio_token_cost through get_model_info Every proxy and router cost lookup goes through get_model_info, which copies cost map keys explicitly, so the new audio cache-read branch always fell back to the text cache-read rate there. Copy the key so models whose audio cache-read rate differs from the text one bill cached audio correctly. --- litellm/types/utils.py | 1 + litellm/utils.py | 1 + .../llm_cost_calc/test_llm_cost_calc_utils.py | 17 +++++++++++++++++ tests/test_litellm/test_utils.py | 8 ++++++++ 4 files changed, 27 insertions(+) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 9eb70e189a1..ef191d79177 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -251,6 +251,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): cache_creation_input_token_cost_priority: float | None # OpenAI priority service tier pricing cache_creation_input_token_cost_ultrafast: ReadOnly[float | None] # OpenAI ultrafast service tier pricing cache_read_input_token_cost: float | None + cache_read_input_audio_token_cost: ReadOnly[float | None] cache_read_input_token_cost_flex: float | None # OpenAI flex service tier pricing cache_read_input_token_cost_priority: float | None # OpenAI priority service tier pricing cache_read_input_token_cost_ultrafast: ReadOnly[float | None] # OpenAI ultrafast service tier pricing diff --git a/litellm/utils.py b/litellm/utils.py index 1a77655a5a4..d048265ecd3 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5866,6 +5866,7 @@ def _get_model_info_helper( "cache_creation_input_token_cost_ultrafast", None ), cache_read_input_token_cost=_model_info.get("cache_read_input_token_cost", None), + cache_read_input_audio_token_cost=_model_info.get("cache_read_input_audio_token_cost", None), prompt_cache_min_tokens=_model_info.get("prompt_cache_min_tokens", None), cache_read_input_token_cost_above_200k_tokens=_model_info.get( "cache_read_input_token_cost_above_200k_tokens", None diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 5d55dfb14a3..5ff1ab62698 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -5235,3 +5235,20 @@ def test_cached_audio_tokens_capped_at_cached_tokens(_local_model_cost_map: None model="gpt-realtime-2", usage=usage, custom_llm_provider="openai" ) assert prompt_cost == pytest.approx(116 * 4e-6 + (167 - 100) * 32e-6 + 100 * 4e-7) + + +def test_cached_audio_tokens_billed_at_audio_cache_rate_through_model_info_lookup(_local_model_cost_map: None) -> None: + usage = Usage( + prompt_tokens=1000, + completion_tokens=0, + total_tokens=1000, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=400, + audio_tokens=600, + cached_tokens=500, + cached_tokens_details={"text_tokens": 100, "audio_tokens": 400}, + ), + ) + + prompt_cost, _ = generic_cost_per_token(model="gpt-realtime-2.1-mini", usage=usage, custom_llm_provider="openai") + assert prompt_cost == pytest.approx(300 * 6e-7 + 100 * 6e-8 + 200 * 1e-5 + 400 * 3e-7) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 19ed31c7b22..d9ca7d9ddee 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -6430,3 +6430,11 @@ def test_completion_finishes_response_metadata_before_handing_the_response_to_th assert snapshot["litellm_call_id"] assert snapshot["response_cost"] is not None assert snapshot["api_base"] + + +def test_get_model_info_carries_cache_read_input_audio_token_cost(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + info = litellm.get_model_info("gpt-realtime-2.1-mini", custom_llm_provider="openai") + assert info["cache_read_input_audio_token_cost"] == 3e-07 + assert info["cache_read_input_token_cost"] == 6e-08 From 305caa8260fb98f821aa69eeba466a162b69c1c4 Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 12 Sep 2026 23:01:41 +0000 Subject: [PATCH 014/112] test: drop unrelated reformatting from merge resolution Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/test_cost_calculator.py | 335 +++++++++++++++------ 1 file changed, 241 insertions(+), 94 deletions(-) diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index f2ee2cdbd9a..c23f5c08a70 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -1,3 +1,4 @@ + import json from pathlib import Path from typing import Final @@ -148,7 +149,9 @@ def test_jina_rerank_bills_total_tokens_at_input_rate_only(_local_model_cost_map def test_cost_calculator_with_response_cost_in_additional_headers(): class MockResponse(BaseModel): - _hidden_params = {"additional_headers": {"llm_provider-x-litellm-response-cost": 1000}} + _hidden_params = { + "additional_headers": {"llm_provider-x-litellm-response-cost": 1000} + } result = response_cost_calculator( response_object=MockResponse(), @@ -204,9 +207,7 @@ def test_vertex_lyria_speech_cost( call_type=call_type, ) - expected: Final = ( - 0 if runtime_state == "custom_zero" else expected_cost * (2 if runtime_state == "custom_price" else 1) - ) + expected: Final = 0 if runtime_state == "custom_zero" else expected_cost * (2 if runtime_state == "custom_price" else 1) assert cost == pytest.approx(expected) @@ -333,12 +334,13 @@ def test_cost_calculator_with_usage(_local_model_cost_map, monkeypatch): # Step 1: Test a model where input_cost_per_image_token is not set. # In this case the calculation should use input_cost_per_token as fallback. - assert model_info.get("input_cost_per_image_token") is None, ( - "Test case expects that input_cost_per_image_token is not set" - ) + assert ( + model_info.get("input_cost_per_image_token") is None + ), "Test case expects that input_cost_per_image_token is not set" expected_cost = ( - usage.prompt_tokens_details.audio_tokens * model_info["input_cost_per_audio_token"] + usage.prompt_tokens_details.audio_tokens + * model_info["input_cost_per_audio_token"] + usage.prompt_tokens_details.text_tokens * model_info["input_cost_per_token"] + usage.prompt_tokens_details.image_tokens * model_info["input_cost_per_token"] + usage.completion_tokens * model_info["output_cost_per_token"] @@ -373,9 +375,12 @@ def test_cost_calculator_with_usage(_local_model_cost_map, monkeypatch): ) expected_cost = ( - usage.prompt_tokens_details.audio_tokens * temp_model_info_object["input_cost_per_audio_token"] - + usage.prompt_tokens_details.text_tokens * temp_model_info_object["input_cost_per_token"] - + usage.prompt_tokens_details.image_tokens * temp_model_info_object["input_cost_per_image_token"] + usage.prompt_tokens_details.audio_tokens + * temp_model_info_object["input_cost_per_audio_token"] + + usage.prompt_tokens_details.text_tokens + * temp_model_info_object["input_cost_per_token"] + + usage.prompt_tokens_details.image_tokens + * temp_model_info_object["input_cost_per_image_token"] + usage.completion_tokens * temp_model_info_object["output_cost_per_token"] ) @@ -385,11 +390,14 @@ def test_cost_calculator_with_usage(_local_model_cost_map, monkeypatch): def test_transcription_cost_uses_token_pricing(_local_model_cost_map): from litellm import completion_cost + usage = Usage( prompt_tokens=14, completion_tokens=45, total_tokens=59, - prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=0, audio_tokens=14), + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=0, audio_tokens=14 + ), ) response = TranscriptionResponse(text="demo text") response.usage = usage @@ -433,6 +441,7 @@ def test_transcription_token_pricing_is_provider_aware(_local_model_cost_map): def test_transcription_cost_falls_back_to_duration(_local_model_cost_map): from litellm import completion_cost + response = TranscriptionResponse(text="demo text") response.duration = 10.0 @@ -453,6 +462,7 @@ def test_vertex_chirp_3_transcription_cost_from_duration(_local_model_cost_map): every transcription priced to $0.00 instead of using input_cost_per_second.""" from litellm import completion_cost + response = TranscriptionResponse(text="demo text") response.duration = 18.0 @@ -476,7 +486,9 @@ def test_handle_realtime_stream_cost_calculation(): {"type": "session.created", "session": {"model": "gpt-3.5-turbo"}}, { "type": "response.done", - "response": {"usage": {"input_tokens": 100, "output_tokens": 50, "total_tokens": 150}}, + "response": { + "usage": {"input_tokens": 100, "output_tokens": 50, "total_tokens": 150} + }, }, { "type": "response.done", @@ -507,7 +519,9 @@ def test_handle_realtime_stream_cost_calculation(): expected_cost = (300 * 0.0015 / 1000) + ( # input tokens (100 + 200) 150 * 0.002 / 1000 ) # output tokens (50 + 100) - assert abs(cost - expected_cost) <= 0.00075 # Allow small floating point differences + assert ( + abs(cost - expected_cost) <= 0.00075 + ) # Allow small floating point differences # Test with different model name in session results[0]["session"]["model"] = "gpt-4" @@ -587,7 +601,14 @@ def test_handle_realtime_stream_cost_calculation_stores_cost_breakdown(): assert logging_obj.cost_breakdown is not None assert logging_obj.cost_breakdown["input_cost"] > 0 assert logging_obj.cost_breakdown["output_cost"] > 0 - assert abs(logging_obj.cost_breakdown["input_cost"] + logging_obj.cost_breakdown["output_cost"] - total_cost) < 1e-9 + assert ( + abs( + logging_obj.cost_breakdown["input_cost"] + + logging_obj.cost_breakdown["output_cost"] + - total_cost + ) + < 1e-9 + ) assert abs(logging_obj.cost_breakdown["total_cost"] - total_cost) < 1e-9 @@ -661,7 +682,9 @@ def test_realtime_logging_object_allows_null_transcript_in_conversation_item_add }, ] - usage = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results(results=results) + usage = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results( + results=results + ) logging_result = RealtimeAPITokenUsageProcessor.create_logging_realtime_object( usage=usage, results=results, @@ -711,7 +734,9 @@ def test_realtime_logging_object_does_not_validate_unknown_event_types(): }, ] - usage = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results(results=results) + usage = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results( + results=results + ) # On unfixed code this raises pydantic ValidationError instead of returning. logging_result = RealtimeAPITokenUsageProcessor.create_logging_realtime_object( usage=usage, @@ -723,7 +748,8 @@ def test_realtime_logging_object_does_not_validate_unknown_event_types(): unknown_types = { r["type"] for r in logging_result.results - if r["type"] in ("rate_limits.updated", "response.function_call_arguments.delta") + if r["type"] + in ("rate_limits.updated", "response.function_call_arguments.delta") } assert unknown_types == { "rate_limits.updated", @@ -756,7 +782,9 @@ def test_realtime_transcription_duration_cost(monkeypatch): "type": "session.created", "session": { "type": "transcription", - "audio": {"input": {"transcription": {"model": "gpt-realtime-whisper"}}}, + "audio": { + "input": {"transcription": {"model": "gpt-realtime-whisper"}} + }, }, }, { @@ -771,7 +799,9 @@ def test_realtime_transcription_duration_cost(monkeypatch): }, ] - combined = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results(results=results) + combined = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results( + results=results + ) logging_obj = Logging( model="gpt-realtime-whisper", messages=[], @@ -864,7 +894,9 @@ def test_realtime_transcription_token_billed_fallback(monkeypatch): # gpt-4o-transcribe: input_cost_per_audio_token = 2.5e-06, input_cost_per_token = 2.5e-06, # output_cost_per_token = 1e-05 - model_info = litellm.get_model_info(model="gpt-4o-transcribe", custom_llm_provider="openai") + model_info = litellm.get_model_info( + model="gpt-4o-transcribe", custom_llm_provider="openai" + ) usage = { "type": "tokens", "input_tokens": 40, @@ -945,7 +977,10 @@ def test_get_transcription_model_falls_back_to_session_model(monkeypatch): mock_response=True, ) - assert result._hidden_params["response_cost"] > result_2._hidden_params["response_cost"] + assert ( + result._hidden_params["response_cost"] + > result_2._hidden_params["response_cost"] + ) model_info = router.get_deployment_model_info( model_id="my-unique-model-id", model_name="anthropic/claude-sonnet-4-5-20250929" @@ -1108,7 +1143,9 @@ def test_tiered_pricing_only_deployment_selects_router_model_id(): assert entry.get("input_cost_per_token") is None assert entry.get("tiered_pricing") is not None # The stripped shared alias must not carry tiered pricing. - assert litellm.model_cost["dashscope/qwen-tier-only-test"].get("tiered_pricing") is None + assert ( + litellm.model_cost["dashscope/qwen-tier-only-test"].get("tiered_pricing") is None + ) selected = _select_model_name_for_cost_calc( model="dashscope/qwen-tier-only-test", @@ -1188,7 +1225,9 @@ def test_azure_realtime_cost_calculator(_local_model_cost_map): combined_usage_object=Usage( prompt_tokens=100, completion_tokens=100, - prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=10, audio_tokens=90), + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=10, audio_tokens=90 + ), ), custom_llm_provider="azure", litellm_model_name="my-custom-azure-deployment", @@ -1207,6 +1246,7 @@ def test_azure_audio_output_cost_calculation(_local_model_cost_map): """ from litellm.types.utils import Choices, CompletionTokensDetailsWrapper, Message + # Scenario from issue #19764: # Input: 17 text tokens, 0 audio tokens # Output: 110 text tokens, 482 audio tokens @@ -1262,10 +1302,14 @@ def test_azure_audio_output_cost_calculation(_local_model_cost_map): wrong_total_cost = expected_input_cost + wrong_output_cost # Verify audio tokens are NOT charged at text rate (the bug) - assert abs(cost - wrong_total_cost) > 0.001, "Bug: Audio tokens are being charged at text token rate" + assert ( + abs(cost - wrong_total_cost) > 0.001 + ), "Bug: Audio tokens are being charged at text token rate" # Verify cost matches - assert abs(cost - expected_total_cost) < 0.0000001, f"Expected cost {expected_total_cost}, got {cost}" + assert ( + abs(cost - expected_total_cost) < 0.0000001 + ), f"Expected cost {expected_total_cost}, got {cost}" def test_default_image_cost_calculator(monkeypatch): @@ -1279,7 +1323,9 @@ def test_default_image_cost_calculator(monkeypatch): monkeypatch.setattr( litellm, "model_cost", - {"azure/bf9001cd7209f5734ecb4ab937a5a0e2ba5f119708bd68f184db362930f9dc7b": temp_object}, + { + "azure/bf9001cd7209f5734ecb4ab937a5a0e2ba5f119708bd68f184db362930f9dc7b": temp_object + }, ) args = { @@ -1495,7 +1541,9 @@ def test_gemini_25_implicit_caching_cost(): expected_cost = 0.00068708 # Allow for small floating point differences - assert abs(result - expected_cost) < 1e-8, f"Expected cost {expected_cost}, but got {result}" + assert ( + abs(result - expected_cost) < 1e-8 + ), f"Expected cost {expected_cost}, but got {result}" print(f"✓ Gemini 2.5 implicit caching cost calculation is correct: ${result:.8f}") @@ -1566,7 +1614,9 @@ def test_log_context_cost_calculation(): # Get model info to understand the pricing from litellm import get_model_info - model_info = get_model_info(model="claude-4-sonnet-20250514", custom_llm_provider="anthropic") + model_info = get_model_info( + model="claude-4-sonnet-20250514", custom_llm_provider="anthropic" + ) # Calculate expected cost based on actual model pricing input_cost_per_token = model_info.get("input_cost_per_token", 0) @@ -1574,8 +1624,12 @@ def test_log_context_cost_calculation(): cache_creation_cost_per_token = model_info.get("cache_creation_input_token_cost", 0) # Check if tiered pricing is applied - input_cost_above_200k = model_info.get("input_cost_per_token_above_200k_tokens", input_cost_per_token) - output_cost_above_200k = model_info.get("output_cost_per_token_above_200k_tokens", output_cost_per_token) + input_cost_above_200k = model_info.get( + "input_cost_per_token_above_200k_tokens", input_cost_per_token + ) + output_cost_above_200k = model_info.get( + "output_cost_per_token_above_200k_tokens", output_cost_per_token + ) cache_creation_above_200k = model_info.get( "cache_creation_input_token_cost_above_200k_tokens", cache_creation_cost_per_token, @@ -1583,23 +1637,31 @@ def test_log_context_cost_calculation(): print(f"DEBUG: Base input cost per token: ${input_cost_per_token:.2e}") print(f"DEBUG: Base output cost per token: ${output_cost_per_token:.2e}") - print(f"DEBUG: Base cache creation cost per token: ${cache_creation_cost_per_token:.2e}") + print( + f"DEBUG: Base cache creation cost per token: ${cache_creation_cost_per_token:.2e}" + ) # Handle tiered pricing - if not available, use base pricing if input_cost_above_200k is not None: - print(f"DEBUG: Tiered input cost per token (>200k): ${input_cost_above_200k:.2e}") + print( + f"DEBUG: Tiered input cost per token (>200k): ${input_cost_above_200k:.2e}" + ) else: print("DEBUG: No tiered input pricing available, using base pricing") input_cost_above_200k = input_cost_per_token if output_cost_above_200k is not None: - print(f"DEBUG: Tiered output cost per token (>200k): ${output_cost_above_200k:.2e}") + print( + f"DEBUG: Tiered output cost per token (>200k): ${output_cost_above_200k:.2e}" + ) else: print("DEBUG: No tiered output pricing available, using base pricing") output_cost_above_200k = output_cost_per_token if cache_creation_above_200k is not None: - print(f"DEBUG: Tiered cache creation cost per token (>200k): ${cache_creation_above_200k:.2e}") + print( + f"DEBUG: Tiered cache creation cost per token (>200k): ${cache_creation_above_200k:.2e}" + ) else: print("DEBUG: No tiered cache creation pricing available, using base pricing") cache_creation_above_200k = cache_creation_cost_per_token @@ -1613,9 +1675,13 @@ def test_log_context_cost_calculation(): print(f"DEBUG: Expected total: ${expected_total:.6f}") # Allow for small floating point differences - assert abs(result - expected_total) < 1e-6, f"Expected cost ${expected_total:.6f}, but got ${result:.6f}" + assert ( + abs(result - expected_total) < 1e-6 + ), f"Expected cost ${expected_total:.6f}, but got ${result:.6f}" - print(f"✓ Log context cost calculation with tiered pricing is correct: ${result:.6f}") + print( + f"✓ Log context cost calculation with tiered pricing is correct: ${result:.6f}" + ) print(f" - Input tokens (300k): ${expected_input_cost:.6f}") print(f" - Output tokens (50k): ${expected_output_cost:.6f}") print(f" - Cache creation (1k): ${expected_cache_cost:.6f}") @@ -1674,7 +1740,8 @@ def test_gemini_25_explicit_caching_cost_direct_usage(): expected_actual_cost = ( model_info["input_cost_per_token"] * usage.prompt_tokens_details.text_tokens - + model_info["cache_read_input_token_cost"] * usage.prompt_tokens_details.cached_tokens + + model_info["cache_read_input_token_cost"] + * usage.prompt_tokens_details.cached_tokens + model_info["output_cost_per_token"] * usage.completion_tokens ) @@ -1698,6 +1765,7 @@ def test_azure_ai_cache_cost_calculation(_local_model_cost_map): from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token from litellm.types.utils import PromptTokensDetailsWrapper, Usage + # Register a custom azure_ai model with cache pricing test_model_id = "test-azure-ai-claude-model" litellm.register_model( @@ -1746,12 +1814,13 @@ def test_azure_ai_cache_cost_calculation(_local_model_cost_map): print(f"Output cost: {output_cost}, Expected: {expected_output_cost}") print(f"Total cost: {total_cost}") - assert abs(input_cost - expected_input_cost) < 1e-10, ( - f"Input cost mismatch: got {input_cost}, expected {expected_input_cost}" - ) - assert abs(output_cost - expected_output_cost) < 1e-10, ( - f"Output cost mismatch: got {output_cost}, expected {expected_output_cost}" - ) + assert ( + abs(input_cost - expected_input_cost) < 1e-10 + ), f"Input cost mismatch: got {input_cost}, expected {expected_input_cost}" + assert ( + abs(output_cost - expected_output_cost) < 1e-10 + ), f"Output cost mismatch: got {output_cost}, expected {expected_output_cost}" + AZURE_GPT_5_6_MAP_KEYS = ( @@ -1820,7 +1889,6 @@ def test_azure_gpt_5_6_rates_match_azure_price_page(_local_model_cost_map, model for key in token_cost_keys: assert entry[key] == pytest.approx(global_entry[key] * 1.1), key - def test_vertex_regional_deployment_costs_uplift_over_global(monkeypatch): """ Regression for https://github.com/BerriAI/litellm/issues/34393: two Vertex @@ -1903,6 +1971,7 @@ def test_cost_discount_vertex_ai(monkeypatch): from litellm import completion_cost from litellm.types.utils import Usage + # Create mock response (use a model that exists in model_prices_and_context_window.json) response = ModelResponse( id="test-id", @@ -1931,6 +2000,7 @@ def test_cost_discount_vertex_ai(monkeypatch): custom_llm_provider="vertex_ai", ) + # Verify discount is applied (5% off means 95% of original cost) expected_cost = cost_without_discount * 0.95 assert cost_with_discount == pytest.approx(expected_cost, rel=1e-9) @@ -1948,6 +2018,7 @@ def test_cost_discount_not_applied_to_other_providers(monkeypatch): from litellm import completion_cost from litellm.types.utils import Usage + # Create mock response for OpenAI response = ModelResponse( id="test-id", @@ -1976,6 +2047,7 @@ def test_cost_discount_not_applied_to_other_providers(monkeypatch): custom_llm_provider="openai", ) + # Costs should be the same (no discount applied to OpenAI) assert cost_with_selective_discount == cost_without_discount @@ -1991,6 +2063,7 @@ def test_cost_margin_percentage(monkeypatch): from litellm import completion_cost from litellm.types.utils import Usage + # Create mock response response = ModelResponse( id="test-id", @@ -2019,6 +2092,7 @@ def test_cost_margin_percentage(monkeypatch): custom_llm_provider="openai", ) + # Verify margin is applied (10% margin means 110% of original cost) expected_cost = cost_without_margin * 1.10 assert cost_with_margin == pytest.approx(expected_cost, rel=1e-9) @@ -2036,6 +2110,7 @@ def test_cost_margin_fixed_amount(monkeypatch): from litellm import completion_cost from litellm.types.utils import Usage + # Create mock response response = ModelResponse( id="test-id", @@ -2064,6 +2139,7 @@ def test_cost_margin_fixed_amount(monkeypatch): custom_llm_provider="openai", ) + # Verify fixed margin is applied expected_cost = cost_without_margin + 0.001 assert cost_with_margin == pytest.approx(expected_cost, rel=1e-9) @@ -2081,6 +2157,7 @@ def test_cost_margin_combined(monkeypatch): from litellm import completion_cost from litellm.types.utils import Usage + # Create mock response response = ModelResponse( id="test-id", @@ -2100,7 +2177,9 @@ def test_cost_margin_combined(monkeypatch): ) # Set 8% margin + $0.0005 fixed for openai - monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"percentage": 0.08, "fixed_amount": 0.0005}}) + monkeypatch.setattr(litellm, "cost_margin_config", { + "openai": {"percentage": 0.08, "fixed_amount": 0.0005} + }) # Calculate cost with margin cost_with_margin = completion_cost( @@ -2109,6 +2188,7 @@ def test_cost_margin_combined(monkeypatch): custom_llm_provider="openai", ) + # Verify combined margin is applied expected_cost = cost_without_margin * 1.08 + 0.0005 assert cost_with_margin == pytest.approx(expected_cost, rel=1e-9) @@ -2126,6 +2206,7 @@ def test_cost_margin_global(monkeypatch): from litellm import completion_cost from litellm.types.utils import Usage + # Create mock response response = ModelResponse( id="test-id", @@ -2154,6 +2235,7 @@ def test_cost_margin_global(monkeypatch): custom_llm_provider="openai", ) + # Verify global margin is applied expected_cost = cost_without_margin * 1.05 assert cost_with_global_margin == pytest.approx(expected_cost, rel=1e-9) @@ -2171,6 +2253,7 @@ def test_cost_margin_provider_overrides_global(monkeypatch): from litellm import completion_cost from litellm.types.utils import Usage + # Create mock response response = ModelResponse( id="test-id", @@ -2199,13 +2282,16 @@ def test_cost_margin_provider_overrides_global(monkeypatch): custom_llm_provider="openai", ) + # Verify provider-specific margin is used (not global) expected_cost = cost_without_margin * 1.10 # 10% from provider, not 5% from global assert cost_with_provider_margin == pytest.approx(expected_cost, rel=1e-9) print("✓ Cost margin provider override test passed:") print(f" - Original cost: ${cost_without_margin:.6f}") - print(f" - Cost with provider margin (10%, overrides 5% global): ${cost_with_provider_margin:.6f}") + print( + f" - Cost with provider margin (10%, overrides 5% global): ${cost_with_provider_margin:.6f}" + ) print(f" - Margin added: ${cost_with_provider_margin - cost_without_margin:.6f}") @@ -2216,6 +2302,7 @@ def test_cost_margin_with_discount(monkeypatch): from litellm import completion_cost from litellm.types.utils import Usage + # Create mock response response = ModelResponse( id="test-id", @@ -2246,6 +2333,7 @@ def test_cost_margin_with_discount(monkeypatch): custom_llm_provider="openai", ) + # Verify: discount applied first, then margin # Base cost -> discount: base * 0.95 -> margin: (base * 0.95) * 1.10 expected_cost = base_cost * 0.95 * 1.10 @@ -2283,7 +2371,9 @@ def test_azure_image_generation_cost_calculator(): size=None, usage=ImageUsage( input_tokens=0, - input_tokens_details=ImageUsageInputTokensDetails(image_tokens=0, text_tokens=0), + input_tokens_details=ImageUsageInputTokensDetails( + image_tokens=0, text_tokens=0 + ), output_tokens=0, total_tokens=0, ), @@ -2313,6 +2403,7 @@ def test_completion_cost_extracts_service_tier_from_response(_local_model_cost_m """Test that completion_cost extracts service_tier from completion_response object.""" from litellm import completion_cost + # Test with gpt-5-nano which has flex pricing model = "gpt-5-nano" @@ -2353,18 +2444,23 @@ def test_completion_cost_extracts_service_tier_from_response(_local_model_cost_m assert flex_cost < standard_cost, "Flex cost should be less than standard cost" flex_ratio = flex_cost / standard_cost - assert 0.45 <= flex_ratio <= 0.55, f"Flex pricing should be ~50% of standard, got {flex_ratio:.2f}" + assert ( + 0.45 <= flex_ratio <= 0.55 + ), f"Flex pricing should be ~50% of standard, got {flex_ratio:.2f}" def test_completion_cost_extracts_service_tier_from_usage(_local_model_cost_map): """Test that completion_cost extracts service_tier from usage object.""" from litellm import completion_cost + # Test with gpt-5-nano which has flex pricing model = "gpt-5-nano" # Create usage object with service_tier - usage_with_service_tier = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + usage_with_service_tier = Usage( + prompt_tokens=1000, completion_tokens=500, total_tokens=1500 + ) # Set service_tier as an attribute on the usage object setattr(usage_with_service_tier, "service_tier", "flex") @@ -2382,7 +2478,9 @@ def test_completion_cost_extracts_service_tier_from_usage(_local_model_cost_map) ) # Create usage object without service_tier - usage_without_service_tier = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + usage_without_service_tier = Usage( + prompt_tokens=1000, completion_tokens=500, total_tokens=1500 + ) # Create ModelResponse with usage without service_tier response_standard = ModelResponse( @@ -2403,13 +2501,16 @@ def test_completion_cost_extracts_service_tier_from_usage(_local_model_cost_map) assert flex_cost < standard_cost, "Flex cost should be less than standard cost" flex_ratio = flex_cost / standard_cost - assert 0.45 <= flex_ratio <= 0.55, f"Flex pricing should be ~50% of standard, got {flex_ratio:.2f}" + assert ( + 0.45 <= flex_ratio <= 0.55 + ), f"Flex pricing should be ~50% of standard, got {flex_ratio:.2f}" def test_completion_cost_service_tier_priority(_local_model_cost_map): """Test that service_tier extraction follows priority: optional_params > completion_response > usage.""" from litellm import completion_cost + # Test with gpt-5-nano which has flex pricing model = "gpt-5-nano" @@ -2458,13 +2559,16 @@ def test_completion_cost_service_tier_priority(_local_model_cost_map): assert cost_from_usage > 0, "Cost from usage should be greater than 0" # Costs should be similar (all using flex) - assert abs(cost_from_params - cost_from_usage) < 1e-6, "Costs from params and usage should be similar (both flex)" + assert ( + abs(cost_from_params - cost_from_usage) < 1e-6 + ), "Costs from params and usage should be similar (both flex)" def test_completion_cost_service_tier_for_bedrock(_local_model_cost_map): """Test that Bedrock cost calculation applies service_tier-specific pricing.""" from litellm import completion_cost + model = "bedrock/us-east-1/test-bedrock-service-tier-cost-model" litellm.register_model( model_cost={ @@ -2520,6 +2624,7 @@ def test_completion_cost_service_tier_for_anthropic(_local_model_cost_map): from litellm import completion_cost from litellm.llms.anthropic.chat.transformation import AnthropicConfig + model = "claude-test-service-tier-cost-model" litellm.register_model( model_cost={ @@ -2572,6 +2677,7 @@ def test_completion_cost_anthropic_auto_tier_uses_served_priority_rate(_local_mo from litellm import completion_cost from litellm.llms.anthropic.chat.transformation import AnthropicConfig + model = "claude-test-auto-tier-cost-model" litellm.register_model( model_cost={ @@ -2665,6 +2771,7 @@ def test_completion_cost_non_string_service_tier_defers_to_served_tier(_local_mo from litellm import completion_cost from litellm.llms.anthropic.chat.transformation import AnthropicConfig + model = "claude-test-non-string-tier-cost-model" litellm.register_model( model_cost={ @@ -2714,6 +2821,7 @@ def test_completion_cost_non_string_response_service_tier_defers_to_served_tier( from litellm import completion_cost from litellm.llms.anthropic.chat.transformation import AnthropicConfig + model = "claude-test-response-non-string-tier-cost-model" litellm.register_model( model_cost={ @@ -2736,7 +2844,9 @@ def test_completion_cost_non_string_response_service_tier_defers_to_served_tier( }, reasoning_content=None, ) - response = ModelResponse(usage=usage, model=model, service_tier={"name": "priority"}) + response = ModelResponse( + usage=usage, model=model, service_tier={"name": "priority"} + ) cost = completion_cost( completion_response=response, @@ -2759,6 +2869,7 @@ def test_completion_cost_non_string_usage_service_tier_prices_standard(_local_mo """ from litellm import completion_cost + model = "claude-test-usage-non-string-tier-cost-model" litellm.register_model( model_cost={ @@ -2805,6 +2916,7 @@ def test_anthropic_cost_per_token_prices_cache_at_served_tier_with_multiplier(_l ) from litellm.types.utils import PromptTokensDetailsWrapper, Usage + model = "claude-test-priority-cache-fast-model" litellm.register_model( model_cost={ @@ -2830,7 +2942,9 @@ def test_anthropic_cost_per_token_prices_cache_at_served_tier_with_multiplier(_l ) usage.speed = "fast" - prompt_cost, completion_cost = anthropic_cost_per_token(model=model, usage=usage, service_tier="priority") + prompt_cost, completion_cost = anthropic_cost_per_token( + model=model, usage=usage, service_tier="priority" + ) expected_prompt = ((1000 - 200) * 6e-6 + 200 * 0.6e-6) * 2 expected_completion = 500 * 30e-6 * 2 @@ -2960,7 +3074,9 @@ def test_anthropic_fast_multiplier_only_on_models_with_fast_mode(_local_model_co "model", ["claude-sonnet-4-6", "claude-mythos-5", "claude-mythos-preview"], ) -def test_anthropic_us_data_residency_uplift_on_claude_4_6_and_later_models(_local_model_cost_map, monkeypatch, model): +def test_anthropic_us_data_residency_uplift_on_claude_4_6_and_later_models( + _local_model_cost_map, monkeypatch, model +): """ Anthropic bills every Claude 4.6+ model served with ``inference_geo="us"`` at 1.1x, and echoes that geo back in the response usage, so each of these real @@ -3025,26 +3141,28 @@ def test_gemini_cache_tokens_details_no_negative_values(): usage = VertexGeminiConfig._calculate_usage(completion_response) # Text tokens should be non-cached text only: 9402 - 9393 = 9 - assert usage.prompt_tokens_details.text_tokens == 9, ( - f"Expected text_tokens=9, got {usage.prompt_tokens_details.text_tokens}" - ) + assert ( + usage.prompt_tokens_details.text_tokens == 9 + ), f"Expected text_tokens=9, got {usage.prompt_tokens_details.text_tokens}" # Image tokens should be non-cached image only: 258 - 258 = 0 - assert usage.prompt_tokens_details.image_tokens == 0, ( - f"Expected image_tokens=0, got {usage.prompt_tokens_details.image_tokens}" - ) + assert ( + usage.prompt_tokens_details.image_tokens == 0 + ), f"Expected image_tokens=0, got {usage.prompt_tokens_details.image_tokens}" # Total cached should match - assert usage.prompt_tokens_details.cached_tokens == 9651, ( - f"Expected cached_tokens=9651, got {usage.prompt_tokens_details.cached_tokens}" - ) + assert ( + usage.prompt_tokens_details.cached_tokens == 9651 + ), f"Expected cached_tokens=9651, got {usage.prompt_tokens_details.cached_tokens}" # MOST IMPORTANT: text_tokens should NEVER be negative - assert usage.prompt_tokens_details.text_tokens >= 0, ( - f"BUG: text_tokens is negative ({usage.prompt_tokens_details.text_tokens})! This was the issue in #18750" - ) + assert ( + usage.prompt_tokens_details.text_tokens >= 0 + ), f"BUG: text_tokens is negative ({usage.prompt_tokens_details.text_tokens})! This was the issue in #18750" - print("✅ Issue #18750 fix verified: text_tokens is correctly calculated and non-negative") + print( + "✅ Issue #18750 fix verified: text_tokens is correctly calculated and non-negative" + ) def test_gemini_without_cache_tokens_details(): @@ -3112,18 +3230,18 @@ def test_gemini_implicit_caching_cost_calculation(): usage = VertexGeminiConfig._calculate_usage(completion_response) # Verify parsing - assert usage.cache_read_input_tokens == 8000, ( - f"cache_read_input_tokens should be 8000, got {usage.cache_read_input_tokens}" - ) - assert usage.prompt_tokens_details.cached_tokens == 8000, ( - f"cached_tokens should be 8000, got {usage.prompt_tokens_details.cached_tokens}" - ) + assert ( + usage.cache_read_input_tokens == 8000 + ), f"cache_read_input_tokens should be 8000, got {usage.cache_read_input_tokens}" + assert ( + usage.prompt_tokens_details.cached_tokens == 8000 + ), f"cached_tokens should be 8000, got {usage.prompt_tokens_details.cached_tokens}" # CRITICAL: text_tokens should be (10000 - 8000) = 2000, NOT 10000 # This is the fix for issue #16341 - assert usage.prompt_tokens_details.text_tokens == 2000, ( - f"text_tokens should be 2000 (10000 - 8000), got {usage.prompt_tokens_details.text_tokens}" - ) + assert ( + usage.prompt_tokens_details.text_tokens == 2000 + ), f"text_tokens should be 2000 (10000 - 8000), got {usage.prompt_tokens_details.text_tokens}" # Verify cost calculation uses cached token pricing response = ModelResponse( @@ -3161,7 +3279,9 @@ def test_gemini_implicit_caching_cost_calculation(): f"Cached tokens may not be using reduced pricing." ) - print("✅ Issue #16341 fix verified: Gemini implicit caching cost calculated correctly") + print( + "✅ Issue #16341 fix verified: Gemini implicit caching cost calculated correctly" + ) def test_additional_costs_only_for_azure_ai(_local_model_cost_map): @@ -3175,6 +3295,7 @@ def test_additional_costs_only_for_azure_ai(_local_model_cost_map): """ from litellm.cost_calculator import _get_additional_costs + # Non-azure_ai providers should return None result = _get_additional_costs( model="gpt-4o", @@ -3317,7 +3438,12 @@ def test_custom_pricing_applies_cache_creation_input_cost_via_prompt_details(): }, ) - expected = (4000 - 1000 - 500) * 0.0000025 + 1000 * 0.00000025 + 500 * 0.000003125 + 100 * 0.000015 + expected = ( + (4000 - 1000 - 500) * 0.0000025 + + 1000 * 0.00000025 + + 500 * 0.000003125 + + 100 * 0.000015 + ) assert cost == pytest.approx(expected) @@ -3362,7 +3488,9 @@ def test_custom_pricing_applies_cache_creation_input_cost_via_cache_write_tokens }, ) - expected_prompt = (4000 - 1000 - 500) * 0.0000025 + 1000 * 0.00000025 + 500 * 0.000003125 + expected_prompt = ( + (4000 - 1000 - 500) * 0.0000025 + 1000 * 0.00000025 + 500 * 0.000003125 + ) expected_completion = 100 * 0.000015 assert prompt_cost == pytest.approx(expected_prompt) @@ -3402,7 +3530,10 @@ def test_extract_cache_read_tokens_zero_when_missing(): assert _extract_cache_read_tokens({}) == 0 assert _extract_cache_read_tokens({"cache_read_input_tokens": None}) == 0 - assert _extract_cache_read_tokens({"prompt_tokens_details": {"cached_tokens": None}}) == 0 + assert ( + _extract_cache_read_tokens({"prompt_tokens_details": {"cached_tokens": None}}) + == 0 + ) def test_extract_cache_creation_tokens_anthropic_top_level(): @@ -3444,7 +3575,12 @@ def test_extract_cache_creation_tokens_zero_when_missing(): assert _extract_cache_creation_tokens({}) == 0 assert _extract_cache_creation_tokens({"cache_creation_input_tokens": None}) == 0 - assert _extract_cache_creation_tokens({"prompt_tokens_details": {"cache_write_tokens": None}}) == 0 + assert ( + _extract_cache_creation_tokens( + {"prompt_tokens_details": {"cache_write_tokens": None}} + ) + == 0 + ) def test_custom_pricing_anthropic_style_cache_tokens_not_double_counted(): @@ -3571,6 +3707,7 @@ def test_completion_cost_logs_reasoning_and_cache_breakdown(_local_model_cost_ma from litellm.litellm_core_utils.litellm_logging import Logging from litellm.types.utils import Choices, CompletionTokensDetailsWrapper, Message + logging_obj = Logging( model="gemini-2.5-flash", messages=[{"role": "user", "content": "Hello"}], @@ -3597,8 +3734,12 @@ def test_completion_cost_logs_reasoning_and_cache_breakdown(_local_model_cost_ma prompt_tokens=209, completion_tokens=3996, total_tokens=4205, - completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=3114, text_tokens=882), - prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=100, text_tokens=109), + completion_tokens_details=CompletionTokensDetailsWrapper( + reasoning_tokens=3114, text_tokens=882 + ), + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=100, text_tokens=109 + ), ), ) @@ -3664,7 +3805,9 @@ def test_completion_cost_logs_the_rates_it_billed_at(monkeypatch): assert rates is not None assert rates.input_cost_per_token == pytest.approx(6e-6) assert rates.cache_read_input_token_cost == pytest.approx(6e-7) - assert logging_obj.cost_breakdown["cache_read_cost"] == pytest.approx(100_000 * rates.cache_read_input_token_cost) + assert logging_obj.cost_breakdown["cache_read_cost"] == pytest.approx( + 100_000 * rates.cache_read_input_token_cost + ) assert logging_obj.cost_breakdown["output_cost"] == pytest.approx(1_000 * rates.output_cost_per_token) @@ -3835,7 +3978,11 @@ def test_completion_cost_bills_interactions_api_response(): cost = completion_cost(completion_response=response, custom_llm_provider="gemini") reasoning_rate = model_info.get("output_cost_per_reasoning_token") or model_info["output_cost_per_token"] - expected = 100 * model_info["input_cost_per_token"] + 50 * model_info["output_cost_per_token"] + 25 * reasoning_rate + expected = ( + 100 * model_info["input_cost_per_token"] + + 50 * model_info["output_cost_per_token"] + + 25 * reasoning_rate + ) assert cost == pytest.approx(expected) assert cost > 0 @@ -4006,9 +4153,7 @@ def test_completion_cost_prices_anthropic_shaped_cache_read_tokens(_local_model_ assert cost == pytest.approx(3 * 4e-6 + 4014 * 4e-7 + 5 * 2e-5, rel=1e-9) -def _together_chat_response( - model: str, prompt_tokens: int, completion_tokens: int, cached_tokens: int -) -> ModelResponse: +def _together_chat_response(model: str, prompt_tokens: int, completion_tokens: int, cached_tokens: int) -> ModelResponse: return ModelResponse( id="chatcmpl-together-cache", choices=[{"finish_reason": "stop", "index": 0, "message": {"content": "acknowledged", "role": "assistant"}}], @@ -4076,8 +4221,6 @@ def test_completion_cost_together_metadata_only_model_still_uses_size_bucket(_lo ) assert cost == pytest.approx((23 + 15) * 8e-07, rel=1e-9) - - def test_select_model_name_strips_unregistered_alias_prefix(_local_model_cost_map): """A router-facing model_name alias containing "/" whose leading segment is NOT a registered provider must not be double-prefixed into a non-existent cost key. @@ -4318,7 +4461,9 @@ def test_every_one_hour_cache_write_rate_is_double_its_input_rate(): """Guard against pasting one model's 1h cache-write price onto another: every provider LiteLLM tracks (Anthropic, Bedrock, Vertex, Azure) publishes the 1h write at 2x input.""" - cost_map = json.loads((Path(__file__).parents[2] / "model_prices_and_context_window.json").read_text()) + cost_map = json.loads( + (Path(__file__).parents[2] / "model_prices_and_context_window.json").read_text() + ) one_hour_prefix = "cache_creation_input_token_cost_above_1hr" deviations = { (name, key): (entry["input_cost_per_token" + key[len(one_hour_prefix) :]], entry[key]) @@ -4524,7 +4669,9 @@ def test_batch_cost_calculator_gpt_6_astra_bills_half_the_standard_rate(_local_m usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) - prompt_cost, completion_cost = batch_cost_calculator(usage=usage, model="gpt-6-astra", custom_llm_provider="openai") + prompt_cost, completion_cost = batch_cost_calculator( + usage=usage, model="gpt-6-astra", custom_llm_provider="openai" + ) assert prompt_cost == pytest.approx(1000 * 5e-6) assert completion_cost == pytest.approx(500 * 2.5e-5) From 59c4cf94397207fa0604d58dac31027b94b8467f Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 12 Sep 2026 23:37:07 +0000 Subject: [PATCH 015/112] refactor(responses): build InputTokensDetails without post-construction mutation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../transformation.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index c75e5d0f5ea..14957a5e5aa 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -2748,17 +2748,20 @@ class LiteLLMCompletionResponsesConfig: cache_write_tokens: Final = getattr(prompt_details, "cache_write_tokens", None) or getattr( prompt_details, "cache_creation_tokens", None ) - input_tokens_details: Final = InputTokensDetails( + cache_write_extra: Final[Mapping[str, int]] = ( + MappingProxyType({"cache_write_tokens": cache_write_tokens}) + if cache_write_tokens is not None + else MappingProxyType({}) + ) + response_usage.input_tokens_details = InputTokensDetails( cached_tokens=prompt_details.cached_tokens if prompt_details.cached_tokens is not None else 0, text_tokens=prompt_details.text_tokens, audio_tokens=prompt_details.audio_tokens, cached_tokens_details=( cached_tokens_details if isinstance(cached_tokens_details, CachedTokensDetails) else None ), + **cache_write_extra, ) - if cache_write_tokens is not None: - setattr(input_tokens_details, "cache_write_tokens", cache_write_tokens) - response_usage.input_tokens_details = input_tokens_details # Translate completion_tokens_details to output_tokens_details if hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details is not None: From 4b84c83788a8e9e4db02b0b85a5f43fe0c49ea30 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 12 Sep 2026 16:59:23 -0700 Subject: [PATCH 016/112] fix(cost): split the cache read breakdown at the audio cache-read rate --- .../litellm_core_utils/llm_cost_calc/utils.py | 32 +++++++++++++++---- .../llm_cost_calc/test_llm_cost_calc_utils.py | 23 +++++++++++++ 2 files changed, 48 insertions(+), 7 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index dc689ca9618..8fc428b38ae 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -1355,6 +1355,7 @@ class BilledTokenRates: input_cost_per_token: float output_cost_per_token: float cache_read_input_token_cost: float + cache_read_input_audio_token_cost: float cache_creation_input_token_cost: float cache_creation_input_token_cost_above_1hr: float output_cost_per_reasoning_token: float @@ -1366,6 +1367,7 @@ class BilledTokenRates: input_cost_per_token=self.input_cost_per_token * multiplier, output_cost_per_token=self.output_cost_per_token * multiplier, cache_read_input_token_cost=self.cache_read_input_token_cost * multiplier, + cache_read_input_audio_token_cost=self.cache_read_input_audio_token_cost * multiplier, cache_creation_input_token_cost=self.cache_creation_input_token_cost * multiplier, cache_creation_input_token_cost_above_1hr=self.cache_creation_input_token_cost_above_1hr * multiplier, output_cost_per_reasoning_token=self.output_cost_per_reasoning_token * multiplier, @@ -1389,15 +1391,16 @@ def _reasoning_token_count(usage: Usage) -> int: return parsed or _coerce_token_count(getattr(usage, "reasoning_tokens", 0)) -def _cache_token_counts(usage: Usage) -> tuple[int, int, CacheCreationTokenDetails | None]: - """(cache read tokens, cache creation tokens, cache creation details): read from prompt_tokens_details - first, then the private top-level counters the Usage constructor mirrors cache tokens onto for - providers/callers that bypass the details.""" +def _cache_token_counts(usage: Usage) -> tuple[int, int, int, CacheCreationTokenDetails | None]: + """(cache read tokens, cached audio tokens, cache creation tokens, cache creation details): read from + prompt_tokens_details first, then the private top-level counters the Usage constructor mirrors cache + tokens onto for providers/callers that bypass the details.""" parsed: Final = parse_prompt_tokens_details(usage) if usage.prompt_tokens_details is not None else None parsed_read: Final = parsed["cache_hit_tokens"] if parsed is not None else 0 parsed_creation: Final = parsed["cache_creation_tokens"] if parsed is not None else 0 return ( parsed_read or _coerce_token_count(getattr(usage, "_cache_read_input_tokens", 0)), + parsed["cache_hit_audio_tokens"] if parsed is not None else 0, parsed_creation or _coerce_token_count(getattr(usage, "_cache_creation_input_tokens", 0)), parsed["cache_creation_token_details"] if parsed is not None else None, ) @@ -1408,11 +1411,13 @@ def _custom_pricing_rates(custom_cost_per_token: CostPerToken) -> BilledTokenRat cache rates (else the input rate) and reasoning at the output rate, as _cost_per_token_custom_pricing_helper does.""" input_rate: Final = custom_cost_per_token["input_cost_per_token"] output_rate: Final = custom_cost_per_token["output_cost_per_token"] + cache_read_rate: Final = custom_cost_per_token.get("cache_read_input_token_cost", input_rate) cache_creation_rate: Final = custom_cost_per_token.get("cache_creation_input_token_cost", input_rate) return BilledTokenRates( input_cost_per_token=input_rate, output_cost_per_token=output_rate, - cache_read_input_token_cost=custom_cost_per_token.get("cache_read_input_token_cost", input_rate), + cache_read_input_token_cost=cache_read_rate, + cache_read_input_audio_token_cost=cache_read_rate, cache_creation_input_token_cost=cache_creation_rate, cache_creation_input_token_cost_above_1hr=cache_creation_rate, output_cost_per_reasoning_token=output_rate, @@ -1449,6 +1454,11 @@ def _cost_map_billed_rates( completion_base_cost=completion_base_cost, current_time=billing_time, ) + audio_cache_read_rate: Final = _get_cost_per_unit( + model_info, + _get_service_tier_cost_key("cache_read_input_audio_token_cost", service_tier), + None, + ) multiplier: Final = ( _get_regional_uplift_multiplier(model_info, data_residency) * get_vertex_regional_endpoint_uplift(model_info, vertex_location) @@ -1458,6 +1468,9 @@ def _cost_map_billed_rates( input_cost_per_token=prompt_base_cost, output_cost_per_token=completion_base_cost, cache_read_input_token_cost=cache_read_cost_rate, + cache_read_input_audio_token_cost=( + audio_cache_read_rate if audio_cache_read_rate is not None else cache_read_cost_rate + ), cache_creation_input_token_cost=cache_creation_cost_rate, cache_creation_input_token_cost_above_1hr=cache_creation_cost_above_1hr_rate, output_cost_per_reasoning_token=reasoning_rate, @@ -1530,7 +1543,9 @@ def get_token_type_cost_breakdown( if rates is None: return TokenTypeCostBreakdown(0.0, 0.0, 0.0) - cache_read_tokens, cache_creation_tokens, cache_creation_token_details = _cache_token_counts(usage) + cache_read_tokens, cached_audio_tokens, cache_creation_tokens, cache_creation_token_details = _cache_token_counts( + usage + ) cache_creation_cost: Final = ( float(cache_creation_tokens) * rates.cache_creation_input_token_cost if custom_cost_per_token is not None @@ -1543,7 +1558,10 @@ def get_token_type_cost_breakdown( ) return TokenTypeCostBreakdown( reasoning_cost=float(_reasoning_token_count(usage)) * rates.output_cost_per_reasoning_token, - cache_read_cost=float(cache_read_tokens) * rates.cache_read_input_token_cost, + cache_read_cost=( + float(cache_read_tokens - cached_audio_tokens) * rates.cache_read_input_token_cost + + float(cached_audio_tokens) * rates.cache_read_input_audio_token_cost + ), cache_creation_cost=cache_creation_cost, rates=rates, ) diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 5ff1ab62698..21b96127a74 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -4006,6 +4006,7 @@ def test_billed_token_rates_follow_the_token_tier_the_breakdown_bills_at(monkeyp input_cost_per_token=6e-6, output_cost_per_token=3e-5, cache_read_input_token_cost=6e-7, + cache_read_input_audio_token_cost=6e-7, cache_creation_input_token_cost=7.5e-6, cache_creation_input_token_cost_above_1hr=0.0, output_cost_per_reasoning_token=3e-5, @@ -5252,3 +5253,25 @@ def test_cached_audio_tokens_billed_at_audio_cache_rate_through_model_info_looku prompt_cost, _ = generic_cost_per_token(model="gpt-realtime-2.1-mini", usage=usage, custom_llm_provider="openai") assert prompt_cost == pytest.approx(300 * 6e-7 + 100 * 6e-8 + 200 * 1e-5 + 400 * 3e-7) + + +def test_cache_read_breakdown_splits_cached_audio_at_the_audio_cache_rate(_local_model_cost_map: None) -> None: + usage = Usage( + prompt_tokens=4863, + completion_tokens=1087, + total_tokens=5950, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=1693, + audio_tokens=3170, + cached_tokens=2816, + cached_tokens_details={"text_tokens": 896, "audio_tokens": 1920}, + ), + ) + + breakdown = get_token_type_cost_breakdown(model="gpt-realtime-2.1-mini", custom_llm_provider="openai", usage=usage) + prompt_cost, _ = generic_cost_per_token(model="gpt-realtime-2.1-mini", usage=usage, custom_llm_provider="openai") + + assert breakdown.cache_read_cost == pytest.approx(896 * 6e-8 + 1920 * 3e-7) + assert breakdown.rates is not None + assert breakdown.rates.cache_read_input_audio_token_cost == pytest.approx(3e-7) + assert prompt_cost == pytest.approx((1693 - 896) * 6e-7 + (3170 - 1920) * 1e-5 + breakdown.cache_read_cost) From 8577d63ff5212173aa5b30bdfe8118b666c74179 Mon Sep 17 00:00:00 2001 From: yassin Date: Sun, 13 Sep 2026 00:14:49 +0000 Subject: [PATCH 017/112] fix(proxy): forward provider request id headers on mapped error responses Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_request_processing.py | 10 +++--- .../proxy/test_common_request_processing.py | 35 +++++++++++++++++++ 2 files changed, 39 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index e6ed60ba177..9b8a98ca0e7 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3445,15 +3445,13 @@ class ProxyBaseLLMRequestProcessing: # a failed request reports no timing, matching /v1/chat/completions read_timing_from_logging_obj=False, ) - # Extract headers from exception - check both e.headers and e.response.headers headers = getattr(e, "headers", None) or {} if not headers: - # Try to get headers from e.response.headers (httpx.Response) _response: Final = attribute_of(e, "response") - if _response is not None: - _response_headers: Final = getattr(_response, "headers", None) - if _response_headers: - headers = get_response_headers(dict(_response_headers)) + _response_headers: Final = getattr(_response, "headers", None) if _response is not None else None + _provider_headers: Final = _response_headers or getattr(e, "litellm_response_headers", None) + if _provider_headers: + headers = get_response_headers(dict(_provider_headers)) headers.update(custom_headers) # Call response headers hook for failure diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index efbb5eedad4..45e436f756a 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -8386,6 +8386,41 @@ async def test_handle_llm_api_exception_forwards_provider_headers_on_http_status assert exc_info.value.headers["llm_provider-x-amzn-requestid"] == "req-passthrough-500" +@pytest.mark.asyncio +async def test_handle_llm_api_exception_forwards_litellm_response_headers_when_response_is_synthetic(): + """Exception mapping hands the proxy a mapped error whose ``response`` is a synthetic empty + ``httpx.Response`` and parks the provider's real headers on ``litellm_response_headers``. + The client must still get the provider request id, as it does on a 200. + """ + import httpx + + from litellm.proxy._types import ProxyException, UserAPIKeyAuth + + mapped = litellm.BadRequestError( + message="OpenAIException - max_tokens is too large: 999999999.", + model="gpt-4o-mini", + llm_provider="openai", + ) + mapped.litellm_response_headers = httpx.Headers({"x-request-id": "req_openai_400"}) + assert dict(mapped.response.headers) == {} + + processor = ProxyBaseLLMRequestProcessing(data={}) + proxy_logging_obj = MagicMock() + proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={}) + + with pytest.raises(ProxyException) as exc_info: + await processor._handle_llm_api_exception( + e=mapped, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + proxy_logging_obj=proxy_logging_obj, + ) + + assert exc_info.value.code == "400" + assert "max_tokens is too large: 999999999." in exc_info.value.message + assert exc_info.value.headers["llm_provider-x-request-id"] == "req_openai_400" + + class TestBackgroundResponseRetrievalGovernance: """LIT-7175: retrieving a background Response attaches the model's post_call policy pipelines.""" From f41c8556b5e22d1cd980df9ba2df7c0b8336d545 Mon Sep 17 00:00:00 2001 From: yassin Date: Sun, 13 Sep 2026 00:05:50 +0000 Subject: [PATCH 018/112] feat(jwt): allow virtual_key_claim_field per issuer Multi-IdP deployments can now set virtual_key_claim_field and unregistered_jwt_client_behavior on a JWTIssuerConfig entry. Tokens from that issuer use the issuer-specific claim path and no-match policy for the virtual key mapping lookup; issuers that omit them keep the global values. The auth flow now enters the mapping lookup when any issuer configures the field, not only when the global field is set. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 30 +++ litellm/proxy/auth/user_api_key_auth.py | 19 +- .../proxy/auth/test_user_api_key_auth.py | 225 +++++++++++++++++- tests/test_litellm/proxy/test__types.py | 58 +++++ 4 files changed, 324 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index ae6c042ab3a..3ede847370d 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4837,6 +4837,14 @@ class JWTIssuerConfig(BaseModel): default=None, description="Issuer-specific claim path to normalize into LiteLLM's end-user id.", ) + virtual_key_claim_field: str | None = Field( + default=None, + description="Issuer-specific claim path used for the virtual key mapping lookup. Falls back to the global field.", + ) + unregistered_jwt_client_behavior: UnregisteredJWTClientBehavior | None = Field( + default=None, + description="Issuer-specific policy when the virtual key claim has no mapping. Falls back to the global policy.", + ) model_config = { "extra": "forbid", @@ -5063,6 +5071,28 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): super().__init__(**kwargs) + def get_issuer_config(self, issuer: str | None) -> JWTIssuerConfig | None: + if issuer is None or self.issuers is None: + return None + return next((config for config in self.issuers if config.issuer == issuer), None) + + def is_virtual_key_mapping_configured(self) -> bool: + if self.virtual_key_claim_field is not None: + return True + return any(config.virtual_key_claim_field is not None for config in self.issuers or ()) + + def get_virtual_key_claim_field(self, issuer: str | None) -> str | None: + issuer_config: Final = self.get_issuer_config(issuer) + if issuer_config is not None and issuer_config.virtual_key_claim_field is not None: + return issuer_config.virtual_key_claim_field + return self.virtual_key_claim_field + + def get_unregistered_jwt_client_behavior(self, issuer: str | None) -> UnregisteredJWTClientBehavior: + issuer_config: Final = self.get_issuer_config(issuer) + if issuer_config is not None and issuer_config.unregistered_jwt_client_behavior is not None: + return issuer_config.unregistered_jwt_client_behavior + return self.unregistered_jwt_client_behavior + class PrismaCompatibleUpdateDBModel(TypedDict, total=False): model_name: str diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 9828311112e..f1b373da898 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -987,9 +987,12 @@ async def _resolve_jwt_to_virtual_key( - Raises HTTPException: REJECT policy hit, missing claim under REJECT/AUTO_REGISTER, or other policy violations. """ - virtual_key_claim_field: Final = jwt_handler.litellm_jwtauth.virtual_key_claim_field + raw_issuer: Final = jwt_claims.get(JWTHandler.LITELLM_JWT_ISSUER_CLAIM) + normalized_issuer: Final = raw_issuer if isinstance(raw_issuer, str) else None + virtual_key_claim_field: Final = jwt_handler.litellm_jwtauth.get_virtual_key_claim_field(normalized_issuer) if virtual_key_claim_field is None: return None + behavior: Final = jwt_handler.litellm_jwtauth.get_unregistered_jwt_client_behavior(normalized_issuer) claim_value: Final = get_nested_value( data=jwt_claims, @@ -1006,7 +1009,6 @@ async def _resolve_jwt_to_virtual_key( # simply by presenting a JWT that omits the configured field. For # AUTO_REGISTER there is no stable identity to map without a claim # value, so we deny rather than create a sentinel-keyed record. - behavior = jwt_handler.litellm_jwtauth.unregistered_jwt_client_behavior if behavior in ( UnregisteredJWTClientBehavior.REJECT, UnregisteredJWTClientBehavior.AUTO_REGISTER, @@ -1021,7 +1023,13 @@ async def _resolve_jwt_to_virtual_key( return None cache_key: Final = jwt_key_mapping_cache_key(virtual_key_claim_field, str(claim_value)) - cached_mapping: Final = await user_api_key_cache.async_get_cache(cache_key) + raw_cached_mapping: Final = await user_api_key_cache.async_get_cache(cache_key) + sentinel_written_by_this_policy: Final = behavior == UnregisteredJWTClientBehavior.AUTO_REGISTER + cached_mapping: Final = ( + None + if raw_cached_mapping == _JWT_PROXY_ADMIN_SENTINEL and not sentinel_written_by_this_policy + else raw_cached_mapping + ) if cached_mapping == _JWT_PROXY_ADMIN_SENTINEL: # Previously resolved to a proxy admin via auth_builder; skip the @@ -1030,7 +1038,6 @@ async def _resolve_jwt_to_virtual_key( return None if cached_mapping == "__NO_MAPPING__": - behavior = jwt_handler.litellm_jwtauth.unregistered_jwt_client_behavior if behavior == UnregisteredJWTClientBehavior.REJECT: raise HTTPException( status_code=403, @@ -1093,8 +1100,6 @@ async def _resolve_jwt_to_virtual_key( ) # No mapping found (DB miss or no DB) — apply no-match policy. - behavior = jwt_handler.litellm_jwtauth.unregistered_jwt_client_behavior - if behavior == UnregisteredJWTClientBehavior.REJECT: # Cache the miss before raising so repeated rejections are served from # cache and don't re-query the DB on every request. @@ -1428,7 +1433,7 @@ async def _user_api_key_auth_builder( # unnecessary DB queries in auth_builder do_standard_jwt_auth = True pending_auto_register: _PendingAutoRegister | None = None - if jwt_handler.litellm_jwtauth.virtual_key_claim_field is not None: + if jwt_handler.litellm_jwtauth.is_virtual_key_mapping_configured(): # Decode JWT to get claims without running full auth_builder jwt_claims: dict | None if jwt_handler.litellm_jwtauth.oidc_userinfo_enabled and not is_jwt: diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 0cdbcde6abc..dac6b7aedc5 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -13,7 +13,7 @@ from unittest.mock import ANY, AsyncMock, MagicMock, patch import pytest -from fastapi import status +from fastapi import HTTPException, status import litellm import litellm.proxy.proxy_server @@ -7442,3 +7442,226 @@ async def test_claude_view_never_reinterprets_explicit_names(monkeypatch, layer) assert data["model"] == ("foo" if layer == "unclaimed" else encoded) await _normalize_claude_model(data, token, request, "/v1/messages") assert data["model"] == ("foo" if layer == "unclaimed" else encoded) + + +ISSUER_ONE = "https://issuer-one.example.com" +ISSUER_TWO = "https://issuer-two.example.com" + + +def _per_issuer_virtual_key_jwt_handler( + global_claim_field: str | None, global_behavior: str = "fallback_team_mapping" +) -> MagicMock: + jwt_handler = MagicMock() + jwt_handler.is_jwt.return_value = True + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( + virtual_key_claim_field=global_claim_field, + unregistered_jwt_client_behavior=global_behavior, + issuers=[ + { + "issuer": ISSUER_ONE, + "jwks_url": f"{ISSUER_ONE}/keys", + "audience": "audience-one", + "team_id_jwt_field": "sub", + }, + { + "issuer": ISSUER_TWO, + "jwks_url": f"{ISSUER_TWO}/keys", + "audience": "audience-two", + "virtual_key_claim_field": "sub", + "unregistered_jwt_client_behavior": "reject", + }, + ], + ) + return jwt_handler + + +def _fake_prisma_with_jwt_key_mapping(hashed_token: str | None) -> tuple[SimpleNamespace, AsyncMock]: + find_first = AsyncMock(return_value=None if hashed_token is None else SimpleNamespace(token=hashed_token)) + prisma_client = SimpleNamespace(db=SimpleNamespace(litellm_jwtkeymapping=SimpleNamespace(find_first=find_first))) + return prisma_client, find_first + + +def _mapping_where(claim_name: str, claim_value: str) -> dict[str, str | bool]: + return {"jwt_claim_name": claim_name, "jwt_claim_value": claim_value, "is_active": True} + + +@pytest.mark.asyncio +async def test_per_issuer_virtual_key_claim_field_selects_the_issuer_mapping_for_the_db_lookup(): + from litellm.proxy.auth.user_api_key_auth import _resolve_jwt_to_virtual_key + + jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field=None) + prisma_client, find_first = _fake_prisma_with_jwt_key_mapping("hashed-mapped-key") + user_api_key_cache = DualCache() + await user_api_key_cache.async_set_cache( + key="hashed-mapped-key", + value=UserAPIKeyAuth(token="hashed-mapped-key", api_key="hashed-mapped-key", team_id="svc-team"), + ) + + resolved = await _resolve_jwt_to_virtual_key( + jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_TWO, "sub": "svc-account-7"}, + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + + find_first.assert_awaited_once_with(where=_mapping_where("sub", "svc-account-7")) + assert isinstance(resolved, UserAPIKeyAuth) + assert resolved.token == "hashed-mapped-key" + assert resolved.team_id == "svc-team" + assert await user_api_key_cache.async_get_cache("jwt_key_mapping:sub:svc-account-7") == "hashed-mapped-key" + + +@pytest.mark.asyncio +async def test_per_issuer_reject_behavior_does_not_leak_into_the_team_issuer(): + from litellm.proxy.auth.user_api_key_auth import _resolve_jwt_to_virtual_key + + jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field=None) + prisma_client, find_first = _fake_prisma_with_jwt_key_mapping(None) + + team_issuer_result = await _resolve_jwt_to_virtual_key( + jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_ONE, "sub": "team-alpha"}, + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=DualCache(), + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + assert team_issuer_result is None + find_first.assert_not_awaited() + + with pytest.raises(HTTPException) as exc: + await _resolve_jwt_to_virtual_key( + jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_TWO, "sub": "unknown-svc"}, + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=DualCache(), + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + + assert exc.value.status_code == 403 + assert "No registered mapping for sub='unknown-svc'" in str(exc.value.detail) + find_first.assert_awaited_once_with(where=_mapping_where("sub", "unknown-svc")) + + +@pytest.mark.asyncio +async def test_proxy_admin_sentinel_cached_by_another_issuer_does_not_bypass_reject(): + from litellm.proxy.auth.user_api_key_auth import _JWT_PROXY_ADMIN_SENTINEL, _resolve_jwt_to_virtual_key + + jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field="sub", global_behavior="auto_register") + prisma_client, find_first = _fake_prisma_with_jwt_key_mapping(None) + user_api_key_cache = DualCache() + await user_api_key_cache.async_set_cache(key="jwt_key_mapping:sub:admin-7", value=_JWT_PROXY_ADMIN_SENTINEL) + + auto_register_issuer_result = await _resolve_jwt_to_virtual_key( + jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_ONE, "sub": "admin-7"}, + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + assert auto_register_issuer_result is None + find_first.assert_not_awaited() + + with pytest.raises(HTTPException) as exc: + await _resolve_jwt_to_virtual_key( + jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_TWO, "sub": "admin-7"}, + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + + assert exc.value.status_code == 403 + assert "No registered mapping for sub='admin-7'" in str(exc.value.detail) + find_first.assert_awaited_once_with(where=_mapping_where("sub", "admin-7")) + + +@pytest.mark.asyncio +async def test_issuer_without_virtual_key_claim_field_falls_back_to_the_global_field(): + from litellm.proxy.auth.user_api_key_auth import _resolve_jwt_to_virtual_key + + jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field="client_id") + prisma_client, find_first = _fake_prisma_with_jwt_key_mapping(None) + + with_claim = await _resolve_jwt_to_virtual_key( + jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_ONE, "sub": "team-alpha", "client_id": "app-9"}, + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=DualCache(), + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + without_claim = await _resolve_jwt_to_virtual_key( + jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_ONE, "sub": "team-alpha"}, + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=DualCache(), + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + + assert with_claim is None + assert without_claim is None + find_first.assert_awaited_once_with(where=_mapping_where("client_id", "app-9")) + + +@pytest.mark.asyncio +async def test_auth_flow_enters_virtual_key_mapping_when_only_an_issuer_configures_the_claim_field(): + jwt_token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJzdmMtYWNjb3VudC03In0.signature" + jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field=None) + jwt_handler.auth_jwt = AsyncMock( + return_value={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_TWO, "sub": "svc-account-7"} + ) + mapped_key = UserAPIKeyAuth(token="hashed-mapped-key", api_key="hashed-mapped-key", team_id="svc-team") + + mock_request = MagicMock() + mock_request.url.path = "/v1/chat/completions" + mock_request.method = "POST" + mock_request.headers = {"authorization": f"Bearer {jwt_token}"} + mock_request.query_params = {} + mock_request.state = SimpleNamespace() + + with ( + patch( # test-quality-ok: the builder reads proxy settings from module globals, no injection seam + "litellm.proxy.proxy_server.general_settings", {"enable_jwt_auth": True} + ), + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: module-global proxy state + patch("litellm.proxy.proxy_server.master_key", "sk-master"), # test-quality-ok: module-global proxy state + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), # test-quality-ok: module-global proxy state + patch( # test-quality-ok: module-global proxy state + "litellm.proxy.proxy_server.user_api_key_cache", DualCache() + ), + patch( # test-quality-ok: module-global proxy state + "litellm.proxy.proxy_server.proxy_logging_obj", MagicMock() + ), + patch("litellm.proxy.proxy_server.jwt_handler", jwt_handler), # test-quality-ok: module-global proxy state + patch( # test-quality-ok: the regression is whether the builder reaches this seam at all + "litellm.proxy.auth.user_api_key_auth._resolve_jwt_to_virtual_key", + new_callable=AsyncMock, + return_value=mapped_key, + ) as resolve_mock, + patch( # test-quality-ok: a mapped key must short-circuit standard JWT auth; reaching it is the failure + "litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder", + new_callable=AsyncMock, + side_effect=AssertionError("standard JWT auth must not run for a mapped virtual key"), + ), + ): + result = await _user_api_key_auth_builder( + request=mock_request, + api_key=jwt_token, + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={"model": "gpt-4o-mini"}, + ) + + resolve_mock.assert_awaited_once() + assert resolve_mock.await_args.kwargs["jwt_claims"][JWTHandler.LITELLM_JWT_ISSUER_CLAIM] == ISSUER_TWO + assert result.api_key == "hashed-mapped-key" + assert result.team_id == "svc-team" diff --git a/tests/test_litellm/proxy/test__types.py b/tests/test_litellm/proxy/test__types.py index 26bb1533da4..9e1486ce90f 100644 --- a/tests/test_litellm/proxy/test__types.py +++ b/tests/test_litellm/proxy/test__types.py @@ -277,3 +277,61 @@ def test_team_membership_budget_table_present_still_works(): } result = LiteLLM_TeamMembership.model_validate(data) assert result.litellm_budget_table is None + + +def test_a_jwt_issuer_can_override_the_virtual_key_claim_field_while_other_issuers_keep_the_global_one(): + from litellm.proxy._types import LiteLLM_JWTAuth, UnregisteredJWTClientBehavior + + jwt_auth = LiteLLM_JWTAuth( + virtual_key_claim_field="client_id", + issuers=[ + { + "issuer": "https://team-idp.example.com", + "jwks_url": "https://team-idp.example.com/keys", + "audience": "litellm", + "team_id_jwt_field": "sub", + }, + { + "issuer": "https://service-idp.example.com", + "jwks_url": "https://service-idp.example.com/keys", + "audience": "litellm", + "virtual_key_claim_field": "sub", + "unregistered_jwt_client_behavior": "reject", + }, + ], + ) + + assert jwt_auth.get_virtual_key_claim_field("https://service-idp.example.com") == "sub" + assert jwt_auth.get_unregistered_jwt_client_behavior("https://service-idp.example.com") is ( + UnregisteredJWTClientBehavior.REJECT + ) + assert jwt_auth.get_virtual_key_claim_field("https://team-idp.example.com") == "client_id" + assert jwt_auth.get_unregistered_jwt_client_behavior("https://team-idp.example.com") is ( + UnregisteredJWTClientBehavior.FALLBACK_TEAM_MAPPING + ) + assert jwt_auth.get_virtual_key_claim_field(None) == "client_id" + assert jwt_auth.get_virtual_key_claim_field("https://unknown-idp.example.com") == "client_id" + + +@pytest.mark.parametrize( + ("global_field", "issuer_field", "is_configured"), + ((None, None, False), ("sub", None, True), (None, "sub", True)), +) +def test_virtual_key_mapping_counts_as_configured_when_any_issuer_sets_the_claim_field( + global_field, issuer_field, is_configured +): + from litellm.proxy._types import LiteLLM_JWTAuth + + jwt_auth = LiteLLM_JWTAuth( + virtual_key_claim_field=global_field, + issuers=[ + { + "issuer": "https://idp.example.com", + "jwks_url": "https://idp.example.com/keys", + "audience": "litellm", + "virtual_key_claim_field": issuer_field, + } + ], + ) + + assert jwt_auth.is_virtual_key_mapping_configured() is is_configured From a94c060b84815438dc3a6dda379e0fb7fa60448f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 12 Sep 2026 17:48:00 -0700 Subject: [PATCH 019/112] fix(cost): fill the missing realtime cache-read rates --- ...odel_prices_and_context_window_backup.json | 9 ++++-- model_prices_and_context_window.json | 9 ++++-- .../llm_cost_calc/test_llm_cost_calc_utils.py | 28 +++++++++++++++++++ 3 files changed, 42 insertions(+), 4 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index de9867f9eee..2ce8914ba5a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -5513,7 +5513,8 @@ }, "azure/gpt-realtime-2025-08-28": { "cache_creation_input_audio_token_cost": 4e-06, - "cache_read_input_token_cost": 4e-06, + "cache_read_input_audio_token_cost": 4e-07, + "cache_read_input_token_cost": 4e-07, "deprecation_date": "2027-03-02", "input_cost_per_audio_token": 3.2e-05, "input_cost_per_image_token": 5e-06, @@ -5546,7 +5547,8 @@ }, "azure/gpt-realtime-1.5-2026-02-23": { "cache_creation_input_audio_token_cost": 4e-06, - "cache_read_input_token_cost": 4e-06, + "cache_read_input_audio_token_cost": 4e-07, + "cache_read_input_token_cost": 4e-07, "deprecation_date": "2027-08-24", "input_cost_per_audio_token": 3.2e-05, "input_cost_per_image_token": 5e-06, @@ -5683,6 +5685,7 @@ }, "azure/gpt-realtime-mini": { "cache_creation_input_audio_token_cost": 3e-07, + "cache_read_input_audio_token_cost": 3e-07, "cache_read_input_token_cost": 6e-08, "input_cost_per_audio_token": 1e-05, "input_cost_per_image_token": 8e-07, @@ -5715,6 +5718,7 @@ }, "azure/gpt-realtime-mini-2025-10-06": { "cache_creation_input_audio_token_cost": 3e-07, + "cache_read_input_audio_token_cost": 3e-07, "cache_read_input_token_cost": 6e-08, "input_cost_per_audio_token": 1e-05, "input_cost_per_image_token": 8e-07, @@ -32683,6 +32687,7 @@ "gpt-realtime-mini": { "cache_creation_input_audio_token_cost": 3e-07, "cache_read_input_audio_token_cost": 3e-07, + "cache_read_input_token_cost": 6e-08, "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 1e-05, "input_cost_per_token": 6e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index de9867f9eee..2ce8914ba5a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -5513,7 +5513,8 @@ }, "azure/gpt-realtime-2025-08-28": { "cache_creation_input_audio_token_cost": 4e-06, - "cache_read_input_token_cost": 4e-06, + "cache_read_input_audio_token_cost": 4e-07, + "cache_read_input_token_cost": 4e-07, "deprecation_date": "2027-03-02", "input_cost_per_audio_token": 3.2e-05, "input_cost_per_image_token": 5e-06, @@ -5546,7 +5547,8 @@ }, "azure/gpt-realtime-1.5-2026-02-23": { "cache_creation_input_audio_token_cost": 4e-06, - "cache_read_input_token_cost": 4e-06, + "cache_read_input_audio_token_cost": 4e-07, + "cache_read_input_token_cost": 4e-07, "deprecation_date": "2027-08-24", "input_cost_per_audio_token": 3.2e-05, "input_cost_per_image_token": 5e-06, @@ -5683,6 +5685,7 @@ }, "azure/gpt-realtime-mini": { "cache_creation_input_audio_token_cost": 3e-07, + "cache_read_input_audio_token_cost": 3e-07, "cache_read_input_token_cost": 6e-08, "input_cost_per_audio_token": 1e-05, "input_cost_per_image_token": 8e-07, @@ -5715,6 +5718,7 @@ }, "azure/gpt-realtime-mini-2025-10-06": { "cache_creation_input_audio_token_cost": 3e-07, + "cache_read_input_audio_token_cost": 3e-07, "cache_read_input_token_cost": 6e-08, "input_cost_per_audio_token": 1e-05, "input_cost_per_image_token": 8e-07, @@ -32683,6 +32687,7 @@ "gpt-realtime-mini": { "cache_creation_input_audio_token_cost": 3e-07, "cache_read_input_audio_token_cost": 3e-07, + "cache_read_input_token_cost": 6e-08, "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 1e-05, "input_cost_per_token": 6e-07, diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 21b96127a74..7cd19c65887 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -5275,3 +5275,31 @@ def test_cache_read_breakdown_splits_cached_audio_at_the_audio_cache_rate(_local assert breakdown.rates is not None assert breakdown.rates.cache_read_input_audio_token_cost == pytest.approx(3e-7) assert prompt_cost == pytest.approx((1693 - 896) * 6e-7 + (3170 - 1920) * 1e-5 + breakdown.cache_read_cost) + + +@pytest.mark.parametrize( + ("model", "custom_llm_provider", "expected_prompt_cost"), + ( + pytest.param("azure/gpt-realtime-2025-08-28", "azure", 300 * 4e-6 + 100 * 4e-7 + 200 * 3.2e-5 + 400 * 4e-7, id="azure-gpt-realtime"), + pytest.param("azure/gpt-realtime-1.5-2026-02-23", "azure", 300 * 4e-6 + 100 * 4e-7 + 200 * 3.2e-5 + 400 * 4e-7, id="azure-gpt-realtime-1.5"), + pytest.param("azure/gpt-realtime-mini", "azure", 300 * 6e-7 + 100 * 6e-8 + 200 * 1e-5 + 400 * 3e-7, id="azure-gpt-realtime-mini"), + pytest.param("gpt-realtime-mini", "openai", 300 * 6e-7 + 100 * 6e-8 + 200 * 1e-5 + 400 * 3e-7, id="openai-gpt-realtime-mini"), + ), +) +def test_realtime_models_bill_cached_text_and_audio_at_their_cache_read_rates( + _local_model_cost_map: None, model: str, custom_llm_provider: str, expected_prompt_cost: float +) -> None: + usage = Usage( + prompt_tokens=1000, + completion_tokens=0, + total_tokens=1000, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=400, + audio_tokens=600, + cached_tokens=500, + cached_tokens_details={"text_tokens": 100, "audio_tokens": 400}, + ), + ) + + prompt_cost, _ = generic_cost_per_token(model=model, usage=usage, custom_llm_provider=custom_llm_provider) + assert prompt_cost == pytest.approx(expected_prompt_cost) From 555e321cf170ad2d15a2932fccdea104da1c39b6 Mon Sep 17 00:00:00 2001 From: yassin Date: Sun, 13 Sep 2026 01:05:51 +0000 Subject: [PATCH 020/112] fix(router): record flat retry attempts and cap retries from attempted_retries Router.log_retry used to copy the failed attempt's kwargs and metadata into metadata.previous_models. Nothing downstream read those copies, but they carried client credentials into spend logs and grew the payload on every retry. Each attempt now leaves a flat record (model group, deployment id, exception type and string, attempt number), which drops RETRY_BREADCRUMB_EXCLUDED_KWARGS and the per-retry credential masking. num_retries_per_request was enforced from len(previous_models), which only looked at the metadata bucket and never exceeded four records. The sync and async client wrappers and the Rust lifecycle guard now read attempted_retries from whichever metadata bucket the call carries. Resolves LIT-7505 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/__init__.py | 2 +- litellm/litellm_core_utils/core_helpers.py | 13 +++ litellm/router.py | 45 +++----- litellm/rust_bridge/lifecycle.py | 16 +-- litellm/types/router.py | 8 ++ litellm/utils.py | 18 +-- .../test_router_helper_utils.py | 29 +++-- .../rust_bridge/test_lifecycle.py | 30 +++++ tests/test_litellm/test_router.py | 106 +++++++++++------- tests/test_litellm/test_utils.py | 47 ++++++++ tests/test_litellm_rust/ocr/test_lifecycle.py | 2 +- 11 files changed, 209 insertions(+), 107 deletions(-) create mode 100644 tests/test_litellm/rust_bridge/test_lifecycle.py diff --git a/litellm/__init__.py b/litellm/__init__.py index ccfbf80369f..261457d6889 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -538,7 +538,7 @@ context_window_fallbacks: Optional[List] = None content_policy_fallbacks: Optional[List] = None allowed_fails: int = 3 allow_dynamic_callback_disabling: bool = True -num_retries_per_request: Optional[int] = None # for the request overall (incl. fallbacks + model retries) +num_retries_per_request: Optional[int] = None # cap on Router retries of one model group; resets per fallback hop ####### SECRET MANAGERS ##################### secret_manager_client: Optional[Any] = ( None # list of instantiated key management clients - e.g. azure kv, infisical, etc. diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index aa7d6ca1699..ecd9cdac88b 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -303,6 +303,19 @@ def get_metadata_variable_name_from_kwargs( return "litellm_metadata" if "litellm_metadata" in kwargs else "metadata" +def max_retries_per_request_hit(kwargs: Mapping[str, object], num_retries_per_request: int | None) -> bool: + """ + Whether the Router retry about to run (``attempted_retries`` >= 1 in the metadata bucket) is past the cap + """ + if num_retries_per_request is None: + return False + metadata: Final = kwargs.get(get_metadata_variable_name_from_kwargs(kwargs)) + if not isinstance(metadata, Mapping): + return False + attempted_retries: Final = metadata.get("attempted_retries") + return type(attempted_retries) is int and 0 < attempted_retries and num_retries_per_request <= attempted_retries + + def get_or_create_metadata_bucket( request_data: dict, ) -> tuple[Literal["metadata", "litellm_metadata"], dict]: diff --git a/litellm/router.py b/litellm/router.py index 8865543badd..a46ffa83b05 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -96,7 +96,6 @@ from litellm.litellm_core_utils.request_timeout_resolver import ( from litellm.litellm_core_utils.secret_redaction import redact_string from litellm.litellm_core_utils.sensitive_data_masker import ( SensitiveDataMasker, - mask_credentials_in_payload, mask_sensitive_structure, ) from litellm.litellm_core_utils.token_counter import offload_token_count @@ -242,6 +241,7 @@ from litellm.types.router import ( ModelGroupInfo, OptionalPreCallChecks, PreRoutingStrategy, + RetryAttemptRecord, RetryPolicy, RouterCacheEnum, RouterErrors, @@ -623,20 +623,6 @@ def _replay_live_router_model_cost() -> None: set_live_deployment_replay(_replay_live_router_model_cost) -# Kwargs that carry no signal about the failed attempt, so log_retry drops them from a -# breadcrumb entirely: the request payload, the proxy's snapshot of the inbound request (its body -# aliases the live request metadata, earlier breadcrumbs included, so copying it would nest every -# breadcrumb inside the next one), and the router-internal walk state. Credentials are handled -# separately by mask_credentials_in_payload, which scrubs credential-named values from whatever -# kwargs remain rather than trying to enumerate every credential-bearing key here. -RETRY_BREADCRUMB_EXCLUDED_KWARGS: Final = frozenset( - ( - "messages", - "original_function", - "attempted_targets", - "proxy_server_request", - ) -) RETRY_BREADCRUMB_LIMIT: Final = 4 @@ -8374,31 +8360,28 @@ class Router: def log_retry(self, kwargs: dict, e: Exception) -> dict: """ - When a retry or fallback happens, log the details of the just failed model call - similar to Sentry breadcrumbing + When a retry or fallback happens, record which model group, deployment and attempt just failed and why """ _metadata_var: Final = "litellm_metadata" if "litellm_metadata" in kwargs else "metadata" request_metadata: Final[Mapping[str, object]] = kwargs[_metadata_var] - attempt_kwargs: Final = MappingProxyType( - {k: v for k, v in kwargs.items() if k != _metadata_var and k not in RETRY_BREADCRUMB_EXCLUDED_KWARGS} - ) - attempt_metadata: Final = MappingProxyType( - {k: v for k, v in request_metadata.items() if k != "previous_models"} - ) - previous_model: Final = MappingProxyType( - { - "exception_type": type(e).__name__, - "exception_string": str(e), - **attempt_kwargs, - _metadata_var: attempt_metadata, - } - ) + model_group: Final = kwargs.get("model") + model_info: Final = request_metadata.get("model_info") + deployment_id: Final = model_info.get("id") if isinstance(model_info, Mapping) else None + attempted_retries: Final = request_metadata.get("attempted_retries") + attempt_record: Final[RetryAttemptRecord] = { + "model_group": model_group if isinstance(model_group, str) else None, + "deployment_id": deployment_id if isinstance(deployment_id, str) else None, + "exception_type": type(e).__name__, + "exception_string": str(e), + "attempted_retries": attempted_retries if type(attempted_retries) is int else None, + } earlier_breadcrumbs: Final = request_metadata.get("previous_models") kept_breadcrumbs: Final[tuple[object, ...]] = ( tuple(earlier_breadcrumbs)[-(RETRY_BREADCRUMB_LIMIT - 1) :] if isinstance(earlier_breadcrumbs, (list, tuple)) else () ) - breadcrumbs: Final = (*kept_breadcrumbs, mask_credentials_in_payload(previous_model)) + breadcrumbs: Final = (*kept_breadcrumbs, attempt_record) kwargs[_metadata_var]["previous_models"] = breadcrumbs # rebind-ok: the logging object already holds this dict return kwargs diff --git a/litellm/rust_bridge/lifecycle.py b/litellm/rust_bridge/lifecycle.py index f5e0c1b0fc6..f1cc912129d 100644 --- a/litellm/rust_bridge/lifecycle.py +++ b/litellm/rust_bridge/lifecycle.py @@ -99,23 +99,13 @@ def setup( def check_limits(kwargs: Mapping[str, object]) -> None: import litellm + from litellm.litellm_core_utils.core_helpers import max_retries_per_request_hit current_cost: Final = litellm._current_cost # pyright: ignore[reportPrivateUsage] # shared SDK budget counter has no public accessor if litellm.max_budget and current_cost > litellm.max_budget: raise litellm.BudgetExceededError(current_cost=current_cost, max_budget=litellm.max_budget) - metadata: Final = kwargs.get("metadata") - if isinstance(metadata, Mapping): - typed_metadata: Final = cast( # cast-ok: runtime Mapping check establishes read-only metadata - Mapping[str, object], metadata - ) - previous: Final = typed_metadata.get("previous_models") - if ( - isinstance(previous, list) - and litellm.num_retries_per_request is not None - and len(cast(list[object], previous)) # cast-ok: runtime list check establishes the retry history - >= litellm.num_retries_per_request - ): - raise RuntimeError("Max retries per request hit!") + if max_retries_per_request_hit(kwargs, litellm.num_retries_per_request): + raise RuntimeError("Max retries per request hit!") def finalize( diff --git a/litellm/types/router.py b/litellm/types/router.py index fc09c40fe08..fecc0e00f99 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -883,6 +883,14 @@ class RouterModelGroupAliasItem(TypedDict): hidden: bool # if 'True', don't return on `.get_model_list` +class RetryAttemptRecord(TypedDict): + model_group: ReadOnly[str | None] + deployment_id: ReadOnly[str | None] + exception_type: ReadOnly[str] + exception_string: ReadOnly[str] + attempted_retries: ReadOnly[int | None] + + VALID_LITELLM_ENVIRONMENTS = [ "development", "staging", diff --git a/litellm/utils.py b/litellm/utils.py index 394ab4b4094..98881e68986 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -81,7 +81,7 @@ from litellm.constants import ( PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO, TOOL_CHOICE_OBJECT_TOKEN_COUNT, ) -from litellm.litellm_core_utils.core_helpers import normalize_drop_params +from litellm.litellm_core_utils.core_helpers import max_retries_per_request_hit, normalize_drop_params from litellm.litellm_core_utils.fallback_generalizations import ( match_capability_generalizations, ) @@ -1509,12 +1509,8 @@ def client(original_function): call_type = original_function.__name__ if _is_async_request(kwargs): # [OPTIONAL] CHECK MAX RETRIES / REQUEST - if litellm.num_retries_per_request is not None: - # check if previous_models passed in as ['litellm_params']['metadata]['previous_models'] - previous_models = (kwargs.get("metadata") or {}).get("previous_models", None) - if previous_models is not None: - if litellm.num_retries_per_request <= len(previous_models): - raise Exception("Max retries per request hit!") + if max_retries_per_request_hit(kwargs, litellm.num_retries_per_request): + raise Exception("Max retries per request hit!") # MODEL CALL result = original_function(*args, **kwargs) @@ -1573,12 +1569,8 @@ def client(original_function): ) # [OPTIONAL] CHECK MAX RETRIES / REQUEST - if litellm.num_retries_per_request is not None: - # check if previous_models passed in as ['litellm_params']['metadata]['previous_models'] - previous_models = (kwargs.get("metadata") or {}).get("previous_models", None) - if previous_models is not None: - if litellm.num_retries_per_request <= len(previous_models): - raise Exception("Max retries per request hit!") + if max_retries_per_request_hit(kwargs, litellm.num_retries_per_request): + raise Exception("Max retries per request hit!") # [OPTIONAL] CHECK CACHE print_verbose( diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index 7dbac243d55..5b06c5fdb01 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -1,3 +1,4 @@ +import json import os import traceback from dotenv import load_dotenv @@ -628,17 +629,29 @@ def test_deployment_callback_respects_cooldown_time(model_list): assert mock_set.call_args.kwargs["time_to_cooldown"] == 0 -def test_log_retry(model_list): - """Test if the '_log_retry' function is working correctly""" - import time - +@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) +def test_log_retry(model_list, metadata_key): + """log_retry appends one flat record per failed attempt and copies neither the request kwargs nor + the request metadata into it""" router = Router(model_list=model_list) new_kwargs = router.log_retry( - kwargs={"metadata": {}}, - e=Exception(), + kwargs={ + "model": "gpt-3.5-turbo", + "api_key": "sk-must-not-be-recorded", + "messages": [{"role": "user", "content": "hi"}], + metadata_key: {"model_info": {"id": "deployment-1"}, "attempted_retries": 2, "user_api_key": "sk-proxy"}, + }, + e=litellm.RateLimitError(message="slow down", llm_provider="openai", model="gpt-3.5-turbo"), ) - assert "metadata" in new_kwargs - assert "previous_models" in new_kwargs["metadata"] + assert json.loads(json.dumps(new_kwargs[metadata_key]["previous_models"])) == [ + { + "model_group": "gpt-3.5-turbo", + "deployment_id": "deployment-1", + "exception_type": "RateLimitError", + "exception_string": "litellm.RateLimitError: slow down", + "attempted_retries": 2, + } + ] def test_update_usage(model_list): diff --git a/tests/test_litellm/rust_bridge/test_lifecycle.py b/tests/test_litellm/rust_bridge/test_lifecycle.py new file mode 100644 index 00000000000..1f0b5591c2b --- /dev/null +++ b/tests/test_litellm/rust_bridge/test_lifecycle.py @@ -0,0 +1,30 @@ +from typing import Final + +import pytest + +import litellm +from litellm.rust_bridge.lifecycle import check_limits + + +@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) +@pytest.mark.parametrize( + "cap, attempted_retries, refused", + [(5, 5, True), (5, 4, False), (0, 0, False), (0, 1, True)], + ids=[ + "cap-above-four-reached", + "cap-above-four-not-reached", + "first-attempt-passes-cap-of-zero", + "cap-of-zero-refuses-first-retry", + ], +) +def test_check_limits_reads_attempted_retries( + monkeypatch: pytest.MonkeyPatch, metadata_key: str, cap: int, attempted_retries: int, refused: bool +) -> None: + monkeypatch.setattr(litellm, "num_retries_per_request", cap) + monkeypatch.setattr(litellm, "max_budget", None) + kwargs: Final = {"model": "mistral/mistral-ocr-latest", metadata_key: {"attempted_retries": attempted_retries}} + if refused: + with pytest.raises(RuntimeError, match="Max retries per request hit!"): + check_limits(kwargs) + else: + check_limits(kwargs) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index f5e9b2091a0..6150ce287ed 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -10607,6 +10607,7 @@ def _cyclic_fallback_router(num_retries=0): "api_key": "sk-fake", "mock_response": "litellm.InternalServerError", }, + "model_info": {"id": f"{group}-deployment"}, } for group in groups ], @@ -10656,28 +10657,37 @@ async def test_cyclic_fallback_graph_does_not_amplify_one_request(): assert sum(len(message) for message in capture.messages) < 5_000 +_FLAT_ATTEMPT_RECORD_KEYS = frozenset( + {"model_group", "deployment_id", "exception_type", "exception_string", "attempted_retries"} +) +_BREADCRUMB_CREDENTIAL_CANARY = "Bearer sk-ant-oat01-RETRY-BREADCRUMB-CANARY-doNotShip" + + @pytest.mark.asyncio -async def test_retry_breadcrumbs_do_not_carry_the_walk_state(): - """log_retry copies every kwarg into previous_models, which reaches spend logs and - logging callbacks. The set of already-attempted groups is router-internal walk state - with no diagnostic value there, and it is the one entry that is not a plain scalar. - A retry has to be configured for the walk state to reach log_retry at all.""" +async def test_retry_records_are_flat_and_name_the_failed_group_on_fallback_hops(): + """Each failed attempt leaves a flat record in previous_models, which reaches spend logs and + logging callbacks. Nothing downstream reads the failed attempt's kwargs or metadata, and copying + them is what carried client credentials and multiplied the payload on every retry. A fallback hop + calls log_retry too, so the record has to name the group that failed, not the one taken next.""" router = _cyclic_fallback_router(num_retries=1) capture = _LogCapture(logging.ERROR) recorder = _FallbackAttemptRecorder() await _drive_cyclic_fallback(router, capture, recorder) - breadcrumbs = [breadcrumb for hop in recorder.breadcrumbs_per_target for breadcrumb in hop] - assert breadcrumbs, "no retry breadcrumbs were recorded" - assert any( - "fallback_depth" in breadcrumb for breadcrumb in breadcrumbs - ), "no breadcrumb carried router walk state, so this test cannot see the leak" - for breadcrumb in breadcrumbs: - assert "attempted_targets" not in breadcrumb - - -_BREADCRUMB_CREDENTIAL_CANARY = "Bearer sk-ant-oat01-RETRY-BREADCRUMB-CANARY-doNotShip" + records = [record for hop in recorder.breadcrumbs_per_target for record in hop] + assert records, "no retry records were recorded" + for record in records: + assert set(record) == _FLAT_ATTEMPT_RECORD_KEYS + assert record["exception_type"] == "InternalServerError" + assert record["deployment_id"] == f"{record['model_group']}-deployment" + group_failed_before_hop = {"group-b": "group-a", "group-c": "group-b", "group-d": "group-c"} + for failed_target, hop_records in zip(recorder.failed_targets, recorder.breadcrumbs_per_target): + groups = [record["model_group"] for record in hop_records] + first_own_attempt = groups.index(failed_target) + assert groups[first_own_attempt - 1] == group_failed_before_hop[failed_target] + assert set(groups[first_own_attempt:]) == {failed_target} + assert [record["attempted_retries"] for record in hop_records[first_own_attempt:]][:2] == [0, 1] @pytest.mark.parametrize( @@ -10703,22 +10713,20 @@ _BREADCRUMB_CREDENTIAL_CANARY = "Bearer sk-ant-oat01-RETRY-BREADCRUMB-CANARY-doN ], ) @pytest.mark.asyncio -async def test_retry_breadcrumbs_never_carry_a_forwarded_credential(container_key, request_kwargs): - """log_retry copies kwargs into previous_models, which reaches spend logs and logging callbacks. - Any of these kwargs can carry a client's forwarded Authorization token or a provider key, and a - breadcrumb has no diagnostic use for the raw secret. A denylist of key names is always one new - credential kwarg behind, so log_retry scrubs credential-named values by pattern instead: the - container still reaches the breadcrumb, but the raw secret never does, whatever key holds it.""" +async def test_retry_records_never_carry_a_forwarded_credential(container_key, request_kwargs): + """previous_models reaches spend logs and logging callbacks. Any request kwarg can carry a client's + forwarded Authorization token or a provider key, so the record must not carry request kwargs at + all: neither the credential-bearing container nor the raw secret, whatever key holds it.""" router = _cyclic_fallback_router(num_retries=1) capture = _LogCapture(logging.ERROR) metadata = {} await _drive_cyclic_fallback(router, capture, metadata=metadata, **request_kwargs) - breadcrumbs = metadata["previous_models"] - assert breadcrumbs, "no retry breadcrumbs were recorded" - dumped = json.dumps(breadcrumbs, default=str) - assert container_key in dumped, "the credential-bearing kwarg never reached the breadcrumb, so this test cannot see the leak" + records = metadata["previous_models"] + assert records, "no retry records were recorded" + dumped = json.dumps(records) + assert container_key not in dumped assert _BREADCRUMB_CREDENTIAL_CANARY not in dumped @@ -10743,7 +10751,7 @@ async def _fail_one_proxy_shaped_request(router, request_marker): shallow copy of the request, so body["metadata"] is the very same dict the router later stamps previous_models onto.""" metadata = {"request_marker": request_marker} - with pytest.raises(litellm.InternalServerError): + with pytest.raises((litellm.InternalServerError, litellm.APIConnectionError)): await router.acompletion( model="broken-group", messages=[{"role": "user", "content": "hi"}], @@ -10769,34 +10777,52 @@ def _nested_breadcrumb_lists(node): @pytest.mark.asyncio -async def test_retry_breadcrumbs_stay_per_request_and_flat_across_failing_requests(): - """Every failed attempt appends a breadcrumb to metadata["previous_models"], and the proxy's +async def test_retry_records_stay_per_request_and_flat_across_failing_requests(): + """Every failed attempt appends a record to metadata["previous_models"], and the proxy's request snapshot aliases that same metadata dict. Kept on the Router and copied wholesale, - each breadcrumb embedded every earlier one from every earlier request, so the breadcrumb + each breadcrumb once embedded every earlier one from every earlier request, so the breadcrumb tree, and with it the debug repr of the kwargs, roughly doubled on each failed attempt until a single-worker proxy spent minutes in the redaction regex and stopped answering.""" router = _always_failing_router(num_retries=2) - breadcrumbs_per_request = [ + records_per_request = [ await _fail_one_proxy_shaped_request(router, f"request-{request_number}") for request_number in range(1, 7) ] - for request_number, breadcrumbs in enumerate(breadcrumbs_per_request, start=1): - assert len(breadcrumbs) == 3, "one initial attempt plus two retries failed, each leaving one breadcrumb" - assert {breadcrumb["metadata"]["request_marker"] for breadcrumb in breadcrumbs} == {f"request-{request_number}"} - for breadcrumb in breadcrumbs: - assert _nested_breadcrumb_lists(breadcrumb) == [] - assert len({len(repr(breadcrumbs)) for breadcrumbs in breadcrumbs_per_request}) == 1 + for records in records_per_request: + assert [record["attempted_retries"] for record in records] == [0, 1, 2] + for record in records: + assert set(record) == _FLAT_ATTEMPT_RECORD_KEYS + assert _nested_breadcrumb_lists(record) == [] + assert len({len(repr(records)) for records in records_per_request}) == 1 @pytest.mark.asyncio -async def test_retry_breadcrumbs_keep_only_the_last_four_attempts(): +async def test_retry_records_keep_only_the_last_four_attempts(): router = _always_failing_router(num_retries=6) - breadcrumbs = await _fail_one_proxy_shaped_request(router, "request-1") + records = await _fail_one_proxy_shaped_request(router, "request-1") - assert len(breadcrumbs) == 4 - assert [breadcrumb["metadata"]["attempted_retries"] for breadcrumb in breadcrumbs] == [3, 4, 5, 6] + assert [record["attempted_retries"] for record in records] == [3, 4, 5, 6] + + +@pytest.mark.asyncio +async def test_num_retries_per_request_stops_retries_at_caps_above_four(monkeypatch): + """The cap used to be read off len(previous_models), which never exceeds four, so any cap above + four was inert. Reading the Router's attempted_retries counter instead lets a cap of five refuse + retries five and six before they reach the deployment.""" + monkeypatch.setattr(litellm, "num_retries_per_request", 5) + router = _always_failing_router(num_retries=6) + + records = await _fail_one_proxy_shaped_request(router, "request-1") + + assert [record["attempted_retries"] for record in records] == [3, 4, 5, 6] + assert ["Max retries per request hit!" in record["exception_string"] for record in records] == [ + False, + False, + True, + True, + ] @pytest.mark.asyncio diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 835e87aff88..523feb2e54a 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4061,6 +4061,53 @@ class TestMetadataNoneHandling: assert metadata == {} +_RETRY_CAP_CASES: Final = ( + pytest.param(5, {"attempted_retries": 5}, True, id="cap-above-four-reached"), + pytest.param(5, {"attempted_retries": 4}, False, id="cap-above-four-not-reached"), + pytest.param(0, {"attempted_retries": 0}, False, id="first-attempt-passes-cap-of-zero"), + pytest.param(0, {"attempted_retries": 1}, True, id="cap-of-zero-refuses-first-retry"), + pytest.param(5, {"previous_models": ("a", "b", "c", "d", "e")}, False, id="breadcrumb-count-is-not-the-cap"), + pytest.param(5, None, False, id="metadata-none"), +) + + +def _capped_completion_kwargs(metadata_key: str, metadata: object) -> dict[str, object]: + return { + "model": "openai/gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}], + "api_key": "sk-fake", + "mock_response": "ok", + metadata_key: metadata, + } + + +@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) +@pytest.mark.parametrize("cap, metadata, refused", _RETRY_CAP_CASES) +def test_num_retries_per_request_reads_attempted_retries_sync(monkeypatch, metadata_key, cap, metadata, refused): + """num_retries_per_request is enforced from the Router's attempted_retries counter in whichever + metadata bucket the call carries, so callers on litellm_metadata and caps above four both work""" + monkeypatch.setattr(litellm, "num_retries_per_request", cap) + kwargs: Final = _capped_completion_kwargs(metadata_key, metadata) + if refused: + with pytest.raises(Exception, match="Max retries per request hit!"): + litellm.completion(**kwargs) + else: + assert litellm.completion(**kwargs).choices[0].message.content == "ok" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) +@pytest.mark.parametrize("cap, metadata, refused", _RETRY_CAP_CASES) +async def test_num_retries_per_request_reads_attempted_retries_async(monkeypatch, metadata_key, cap, metadata, refused): + monkeypatch.setattr(litellm, "num_retries_per_request", cap) + kwargs: Final = _capped_completion_kwargs(metadata_key, metadata) + if refused: + with pytest.raises(Exception, match="Max retries per request hit!"): + await litellm.acompletion(**kwargs) + else: + assert (await litellm.acompletion(**kwargs)).choices[0].message.content == "ok" + + class TestValidateAndFixThinkingParam: """Tests for validate_and_fix_thinking_param.""" diff --git a/tests/test_litellm_rust/ocr/test_lifecycle.py b/tests/test_litellm_rust/ocr/test_lifecycle.py index e7ebc5b3018..77d9ef167d0 100644 --- a/tests/test_litellm_rust/ocr/test_lifecycle.py +++ b/tests/test_litellm_rust/ocr/test_lifecycle.py @@ -806,7 +806,7 @@ async def test_shared_call_limits_still_reject_before_reading_ocr_file( monkeypatch.setattr(litellm, "_current_cost", 2) monkeypatch.setattr(litellm, "num_retries_per_request", 1 if limit == "retries" else None) expected: Final = litellm.BudgetExceededError if limit == "budget" else RuntimeError - arguments: Final = {"document": {"type": "file", "file": File()}, "metadata": {"previous_models": ["earlier"]}} + arguments: Final = {"document": {"type": "file", "file": File()}, "metadata": {"attempted_retries": 1}} with pytest.raises(expected, match=r"Budget has been exceeded|Max retries per request hit"): await call_aocr(ocr_server, **arguments) if asynchronous else call_ocr(ocr_server, **arguments) assert reads == [] From 566da87771b65665d9d8489898a86db693b5ca06 Mon Sep 17 00:00:00 2001 From: yassin Date: Sun, 13 Sep 2026 01:20:46 +0000 Subject: [PATCH 021/112] test(router): expect the exact error per retry-cap case and drop explanatory docstrings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/core_helpers.py | 3 --- tests/test_litellm/test_router.py | 9 +++------ tests/test_litellm/test_utils.py | 2 -- 3 files changed, 3 insertions(+), 11 deletions(-) diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index ecd9cdac88b..6e76bf9d49e 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -304,9 +304,6 @@ def get_metadata_variable_name_from_kwargs( def max_retries_per_request_hit(kwargs: Mapping[str, object], num_retries_per_request: int | None) -> bool: - """ - Whether the Router retry about to run (``attempted_retries`` >= 1 in the metadata bucket) is past the cap - """ if num_retries_per_request is None: return False metadata: Final = kwargs.get(get_metadata_variable_name_from_kwargs(kwargs)) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 6150ce287ed..eb29f717a20 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -10746,12 +10746,12 @@ def _always_failing_router(num_retries): ) -async def _fail_one_proxy_shaped_request(router, request_marker): +async def _fail_one_proxy_shaped_request(router, request_marker, expected_error=litellm.InternalServerError): """The proxy hands the router a metadata dict and a proxy_server_request whose body is a shallow copy of the request, so body["metadata"] is the very same dict the router later stamps previous_models onto.""" metadata = {"request_marker": request_marker} - with pytest.raises((litellm.InternalServerError, litellm.APIConnectionError)): + with pytest.raises(expected_error): await router.acompletion( model="broken-group", messages=[{"role": "user", "content": "hi"}], @@ -10808,13 +10808,10 @@ async def test_retry_records_keep_only_the_last_four_attempts(): @pytest.mark.asyncio async def test_num_retries_per_request_stops_retries_at_caps_above_four(monkeypatch): - """The cap used to be read off len(previous_models), which never exceeds four, so any cap above - four was inert. Reading the Router's attempted_retries counter instead lets a cap of five refuse - retries five and six before they reach the deployment.""" monkeypatch.setattr(litellm, "num_retries_per_request", 5) router = _always_failing_router(num_retries=6) - records = await _fail_one_proxy_shaped_request(router, "request-1") + records = await _fail_one_proxy_shaped_request(router, "request-1", expected_error=litellm.APIConnectionError) assert [record["attempted_retries"] for record in records] == [3, 4, 5, 6] assert ["Max retries per request hit!" in record["exception_string"] for record in records] == [ diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 523feb2e54a..3f3b8ce5343 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4084,8 +4084,6 @@ def _capped_completion_kwargs(metadata_key: str, metadata: object) -> dict[str, @pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) @pytest.mark.parametrize("cap, metadata, refused", _RETRY_CAP_CASES) def test_num_retries_per_request_reads_attempted_retries_sync(monkeypatch, metadata_key, cap, metadata, refused): - """num_retries_per_request is enforced from the Router's attempted_retries counter in whichever - metadata bucket the call carries, so callers on litellm_metadata and caps above four both work""" monkeypatch.setattr(litellm, "num_retries_per_request", cap) kwargs: Final = _capped_completion_kwargs(metadata_key, metadata) if refused: From 4c022a3089cbfbc377cc700db7ade6309f3b1e1f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 12 Sep 2026 23:34:12 -0700 Subject: [PATCH 022/112] feat(pricing): add azure gpt-chat-latest global and data zone rates --- ...odel_prices_and_context_window_backup.json | 111 ++++++++++++++++++ model_prices_and_context_window.json | 111 ++++++++++++++++++ .../llm_cost_calc/test_llm_cost_calc_utils.py | 31 +++++ .../test_litellm/test_model_prices_schema.py | 8 ++ 4 files changed, 261 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 2220d0e1fe5..86605d77aab 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -7409,6 +7409,80 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure/gpt-chat-latest": { + "cache_read_input_token_cost": 5e-07, + "deprecation_date": "2026-12-02", + "input_cost_per_token": 5e-06, + "litellm_provider": "azure", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "reasoning_effort_levels": [ + "medium" + ], + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, + "azure/chat-latest": { + "cache_read_input_token_cost": 5e-07, + "deprecation_date": "2026-12-02", + "input_cost_per_token": 5e-06, + "litellm_provider": "azure", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "reasoning_effort_levels": [ + "medium" + ], + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, "azure/us/gpt-5.6": { "cache_creation_input_token_cost": 6.875e-06, "cache_creation_input_token_cost_above_272k_tokens": 1.375e-05, @@ -7675,6 +7749,43 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure/us/gpt-chat-latest": { + "cache_read_input_token_cost": 5.5e-07, + "deprecation_date": "2026-12-02", + "input_cost_per_token": 5.5e-06, + "litellm_provider": "azure", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3.3e-05, + "reasoning_effort_levels": [ + "medium" + ], + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, "azure/eu/gpt-5.6": { "cache_creation_input_token_cost": 6.875e-06, "cache_creation_input_token_cost_above_272k_tokens": 1.375e-05, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 2220d0e1fe5..86605d77aab 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -7409,6 +7409,80 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure/gpt-chat-latest": { + "cache_read_input_token_cost": 5e-07, + "deprecation_date": "2026-12-02", + "input_cost_per_token": 5e-06, + "litellm_provider": "azure", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "reasoning_effort_levels": [ + "medium" + ], + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, + "azure/chat-latest": { + "cache_read_input_token_cost": 5e-07, + "deprecation_date": "2026-12-02", + "input_cost_per_token": 5e-06, + "litellm_provider": "azure", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "reasoning_effort_levels": [ + "medium" + ], + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, "azure/us/gpt-5.6": { "cache_creation_input_token_cost": 6.875e-06, "cache_creation_input_token_cost_above_272k_tokens": 1.375e-05, @@ -7675,6 +7749,43 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure/us/gpt-chat-latest": { + "cache_read_input_token_cost": 5.5e-07, + "deprecation_date": "2026-12-02", + "input_cost_per_token": 5.5e-06, + "litellm_provider": "azure", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3.3e-05, + "reasoning_effort_levels": [ + "medium" + ], + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, "azure/eu/gpt-5.6": { "cache_creation_input_token_cost": 6.875e-06, "cache_creation_input_token_cost_above_272k_tokens": 1.375e-05, diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index cbe6fe198c9..3c6a13a7f01 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -2101,6 +2101,37 @@ def test_generic_cost_per_token_azure_gpt_6_astra_foundry_price_sheet( assert completion_cost == pytest.approx(zone_multiplier * output_multiplier * completion_tokens * 5e-5) +@pytest.mark.parametrize( + "model,input_rate,cache_read_rate,output_rate", + [ + ("azure/gpt-chat-latest", 5e-6, 5e-7, 3e-5), + ("azure/chat-latest", 5e-6, 5e-7, 3e-5), + ("azure/us/gpt-chat-latest", 5.5e-6, 5.5e-7, 3.3e-5), + ], +) +def test_generic_cost_per_token_azure_gpt_chat_latest_price_sheet( + _local_model_cost_map, model, input_rate, cache_read_rate, output_rate +): + """The Azure OpenAI price sheet lists GPT-Chat Latest at $5 input, $0.50 cached input and $30 output per 1M + tokens on Global, and $5.50, $0.55 and $33 on Data Zone. Foundry names the product gpt-chat-latest and the + OpenAI API names the same model chat-latest, so both spellings bill the Global sheet. + """ + prompt_tokens = 100000 + cached_tokens = 40000 + completion_tokens = 1000 + usage = Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=cached_tokens), + ) + + prompt_cost, completion_cost = generic_cost_per_token(model=model, usage=usage, custom_llm_provider="azure") + + assert prompt_cost == pytest.approx((prompt_tokens - cached_tokens) * input_rate + cached_tokens * cache_read_rate) + assert completion_cost == pytest.approx(completion_tokens * output_rate) + + def test_generic_cost_per_token_azure_ai_gpt_6_astra_flex_bills_the_standard_rate(_local_model_cost_map): usage = Usage(prompt_tokens=1000, completion_tokens=100, total_tokens=1100) diff --git a/tests/test_litellm/test_model_prices_schema.py b/tests/test_litellm/test_model_prices_schema.py index 0b9dbd23097..e562797fbe8 100644 --- a/tests/test_litellm/test_model_prices_schema.py +++ b/tests/test_litellm/test_model_prices_schema.py @@ -221,6 +221,14 @@ def test_chat_latest_declares_the_one_effort_openai_accepts(prices: dict): assert resolve_supported_reasoning_efforts(prices["chat-latest"], deployment_is_mapped=True) == ("medium",) +@pytest.mark.parametrize("key", ["azure/gpt-chat-latest", "azure/chat-latest", "azure/us/gpt-chat-latest"]) +def test_azure_gpt_chat_latest_declares_the_one_effort_azure_accepts(prices: dict, key: str): + """Azure answers every reasoning_effort on a gpt-chat-latest deployment except medium with + "Unsupported value ... Supported values are: 'medium'", the same fixed level OpenAI's chat-latest + carries, so the Foundry product name and the OpenAI API name both declare that one level.""" + assert resolve_supported_reasoning_efforts(prices[key], deployment_is_mapped=True) == ("medium",) + + BEDROCK_OPENAI_GPT_MARKERS: Final = ("openai.gpt-5.4", "openai.gpt-5.5", "openai.gpt-5.6", "openai.gpt-6-astra") BEDROCK_PROVIDERS: Final = frozenset(("bedrock", "bedrock_converse", "bedrock_mantle")) BEDROCK_ROW_PREFIXES: Final = ("bedrock_mantle/", "us.", "global.") From 76ae35dfcdefbb6b09cc576909c69b54bae5999d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 01:14:35 -0700 Subject: [PATCH 023/112] fix(types): break the CachedTokensDetails import cycle CodeQL flagged two module-level cyclic imports introduced by defining CachedTokensDetails in litellm.types.llms.openai and importing it from litellm.types.utils and litellm.cost_calculator. The class now lives in litellm.types.llms.base, which imports nothing from litellm, and every user imports it from there. Also pins that combining realtime usages where only one response.done carries cached_tokens_details keeps the earlier modality split in both orders, and commits the regenerated dashboard API types. --- litellm/cost_calculator.py | 2 +- .../transformation.py | 2 +- litellm/types/llms/base.py | 6 +++ litellm/types/llms/openai.py | 8 +--- litellm/types/utils.py | 2 +- tests/test_litellm/test_cost_calculator.py | 43 +++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 - 7 files changed, 53 insertions(+), 12 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index ce5c84907a1..f5319776213 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -97,6 +97,7 @@ from litellm.llms.vertex_ai.cost_calculator import cost_router as google_cost_ro from litellm.llms.xai.cost_calculator import cost_per_token as xai_cost_per_token from litellm.responses.utils import ResponseAPILoggingUtils from litellm.types.agents import LiteLLMSendMessageResponse +from litellm.types.llms.base import CachedTokensDetails from litellm.types.llms.openai import ( HttpxBinaryResponseContent, ImageGenerationRequestQuality, @@ -109,7 +110,6 @@ from litellm.types.llms.openai import ( ) from litellm.types.rerank import RerankBilledUnits, RerankResponse from litellm.types.utils import ( - CachedTokensDetails, CallTypesLiteral, LiteLLMRealtimeStreamLoggingObject, LlmProviders, diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 7a215cdbb12..64324c6cad8 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -43,9 +43,9 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.responses.litellm_completion_transformation.session_handler import ( ResponsesSessionHandler, ) +from litellm.types.llms.base import CachedTokensDetails from litellm.types.llms.openai import ( AllMessageValues, - CachedTokensDetails, ChatCompletionAssistantMessage, ChatCompletionImageObject, ChatCompletionImageUrlObject, diff --git a/litellm/types/llms/base.py b/litellm/types/llms/base.py index f09727ad92b..938aa8064c9 100644 --- a/litellm/types/llms/base.py +++ b/litellm/types/llms/base.py @@ -75,3 +75,9 @@ class HiddenParams(OpenAIObject): data: Final = super().model_dump(**kwargs) data["_response_ms"] = self._response_ms return data + + +class CachedTokensDetails(BaseModel): + text_tokens: int | None = None + audio_tokens: int | None = None + image_tokens: int | None = None diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 7a86c27efae..9bdad700d4b 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -82,7 +82,7 @@ from typing_extensions import ( override, ) -from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject +from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject, CachedTokensDetails from litellm.types.responses.main import ( CustomToolCallOutputItem, GenericResponseOutputItem, @@ -1285,12 +1285,6 @@ class OutputTokensDetails(BaseLiteLLMOpenAIResponseObject): model_config = {"extra": "allow"} -class CachedTokensDetails(BaseModel): - text_tokens: int | None = None - audio_tokens: int | None = None - image_tokens: int | None = None - - class InputTokensDetails(BaseLiteLLMOpenAIResponseObject): audio_tokens: int | None = None cached_tokens: int = 0 diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 07e6fc5838b..1d73542c9bb 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -48,6 +48,7 @@ from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.types.llms.base import ( BaseLiteLLMOpenAIResponseObject, + CachedTokensDetails, LiteLLMPydanticObjectBase, ) from litellm.types.mcp import MCPServerCostInfo @@ -60,7 +61,6 @@ from .llms.base import HiddenParams from .llms.openai import ( AllMessageValues, Batch, - CachedTokensDetails, ChatCompletionAnnotation, ChatCompletionReasoningItem, ChatCompletionRedactedThinkingBlock, diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index a16f1b8fe4a..68e9b6143a0 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -19,6 +19,7 @@ from litellm.cost_calculator import ( ) from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse, OCRUsageInfo +from litellm.types.llms.base import CachedTokensDetails from litellm.types.llms.openai import OpenAIRealtimeStreamList from litellm.types.rerank import RerankResponse from litellm.types.utils import ( @@ -4896,6 +4897,48 @@ def test_realtime_combine_sums_nested_cached_tokens_details(): assert combined.prompt_tokens_details.cached_tokens_details.image_tokens is None +@pytest.mark.parametrize("details_first", [True, False]) +def test_realtime_combine_keeps_cached_split_when_only_one_usage_has_details(details_first: bool): + with_details: Final = { + "type": "response.done", + "response": { + "usage": { + "input_tokens": 283, + "output_tokens": 0, + "total_tokens": 283, + "input_token_details": { + "text_tokens": 116, + "audio_tokens": 167, + "cached_tokens": 192, + "cached_tokens_details": {"text_tokens": 64, "audio_tokens": 128}, + }, + } + }, + } + without_details: Final = { + "type": "response.done", + "response": { + "usage": { + "input_tokens": 150, + "output_tokens": 0, + "total_tokens": 150, + "input_token_details": {"text_tokens": 50, "audio_tokens": 100, "cached_tokens": 100}, + } + }, + } + results: OpenAIRealtimeStreamList = ( + [with_details, without_details] if details_first else [without_details, with_details] + ) + + combined = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results( + results=results, + ) + + assert combined.prompt_tokens_details is not None + assert combined.prompt_tokens_details.cached_tokens == 292 + assert combined.prompt_tokens_details.cached_tokens_details == CachedTokensDetails(text_tokens=64, audio_tokens=128) + + def test_usage_without_cached_tokens_details_omits_key(): usage = Usage( prompt_tokens=10, diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 7eadaa6c991..839aa52fa84 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -16781,7 +16781,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) @@ -16887,7 +16886,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) From 6f7882db34b4d15a0dec402b6218c59afc58c012 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 01:54:12 -0700 Subject: [PATCH 024/112] fix(types): import CachedTokensDetails on its own line in openai.py CodeQL resolves `from openai import Omit` in litellm/types/llms/openai.py to the module itself, so every importer of a name whose definition line is in the diff is reported as an unsafe cyclic import. 76ae35dfcd edited the line that defines BaseLiteLLMOpenAIResponseObject there and got two alerts at files this PR does not touch. That line is now byte-identical to main and CachedTokensDetails arrives through a relative import isort keeps separate. --- litellm/types/llms/openai.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 9bdad700d4b..e3eac9b9205 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -82,7 +82,7 @@ from typing_extensions import ( override, ) -from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject, CachedTokensDetails +from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject from litellm.types.responses.main import ( CustomToolCallOutputItem, GenericResponseOutputItem, @@ -91,6 +91,8 @@ from litellm.types.responses.main import ( OutputImageGenerationCall, ) +from .base import CachedTokensDetails + FileContent = IO[bytes] | bytes | PathLike FileTypes = ( From 76b26e41abfe6e72dd846e7272dcb45069b98b73 Mon Sep 17 00:00:00 2001 From: yassin Date: Sun, 13 Sep 2026 09:07:53 +0000 Subject: [PATCH 025/112] fix(router): cool down team deployments on 429 when a sibling serves the same public model Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 10 ++++ litellm/router_utils/cooldown_handlers.py | 5 +- .../router_utils/test_cooldown_handlers.py | 50 +++++++++++++++++++ 3 files changed, 63 insertions(+), 2 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 8865543badd..4929b17f7fc 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1553,6 +1553,16 @@ class Router: return False return sum(len(self.model_name_to_deployment_indices.get(member) or ()) for member in group.models) > 1 + def team_model_has_alternatives(self, deployment_id: str) -> bool: + deployment: Final = self.get_deployment(model_id=deployment_id) + if deployment is None: + return False + team_id: Final = deployment.model_info.team_id + public_model_name: Final = deployment.model_info.team_public_model_name + if team_id is None or public_model_name is None: + return False + return len(self.team_model_to_deployment_indices.get((team_id, public_model_name)) or ()) > 1 + _OVERRIDABLE_ROUTING_STRATEGIES: frozenset[str] = frozenset({"simple-shuffle", *_DEFAULT_SELECTOR_ATTR_BY_STRATEGY}) def _get_request_routing_strategy_override(self, request_kwargs: dict | None) -> str | None: diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index f722b6fd20c..027f0a9ca05 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -343,8 +343,9 @@ def _should_cooldown_deployment( model_group: Final = litellm_router_instance.get_model_group(id=deployment) is_single_deployment_model_group = False if model_group is not None and len(model_group) == 1: - is_single_deployment_model_group = not litellm_router_instance.routing_group_has_alternatives( - requested_model_group + is_single_deployment_model_group = not ( + litellm_router_instance.routing_group_has_alternatives(requested_model_group) + or litellm_router_instance.team_model_has_alternatives(deployment) ) ## CHECK DEPLOYMENT-LEVEL POLICY FIRST (overrides router-level) diff --git a/tests/test_litellm/router_utils/test_cooldown_handlers.py b/tests/test_litellm/router_utils/test_cooldown_handlers.py index 7ee0ed3701b..6f66bb863cc 100644 --- a/tests/test_litellm/router_utils/test_cooldown_handlers.py +++ b/tests/test_litellm/router_utils/test_cooldown_handlers.py @@ -437,3 +437,53 @@ class TestRoutingGroupCooldownAlternatives: ) is False ) + + +class TestTeamModelCooldownAlternatives: + def _router(self, team_deployments: int): + from litellm import Router + + return Router( + model_list=[ + { + "model_name": f"model_name_team-1_{i}", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}, + "model_info": { + "id": f"team-deploy-{i}", + "team_id": "team-1", + "team_public_model_name": "team-gpt-4o-mini", + }, + } + for i in range(team_deployments) + ] + ) + + def test_429_on_team_deployment_with_sibling_cools_down(self): + from litellm.router_utils.cooldown_handlers import _should_cooldown_deployment + + router = self._router(team_deployments=2) + assert ( + _should_cooldown_deployment( + litellm_router_instance=router, + deployment="team-deploy-0", + exception_status=429, + original_exception=Exception("rate limited"), + requested_model_group="team-gpt-4o-mini", + ) + is True + ) + + def test_429_on_only_team_deployment_keeps_single_deployment_exemption(self): + from litellm.router_utils.cooldown_handlers import _should_cooldown_deployment + + router = self._router(team_deployments=1) + assert ( + _should_cooldown_deployment( + litellm_router_instance=router, + deployment="team-deploy-0", + exception_status=429, + original_exception=Exception("rate limited"), + requested_model_group="team-gpt-4o-mini", + ) + is False + ) From 330ba7cbf91d4db59f3e1b433aa937fcdaddbb2d Mon Sep 17 00:00:00 2001 From: yassin Date: Sun, 13 Sep 2026 09:21:09 +0000 Subject: [PATCH 026/112] fix(ui): show the team alias on the model info page and in its raw JSON Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../src/components/ModelInfoEditForm.tsx | 10 ++- .../src/components/model_info_view.test.tsx | 78 +++++++++++++++++++ .../src/components/model_info_view.tsx | 11 ++- 3 files changed, 96 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/src/components/ModelInfoEditForm.tsx b/ui/litellm-dashboard/src/components/ModelInfoEditForm.tsx index d56e65237eb..fb4d90b5ea2 100644 --- a/ui/litellm-dashboard/src/components/ModelInfoEditForm.tsx +++ b/ui/litellm-dashboard/src/components/ModelInfoEditForm.tsx @@ -271,6 +271,7 @@ const displayCost = (localModelData: any, field: TouchedPricingField): string => interface ModelInfoEditFormProps { localModelData: any; modelData: { model_info: { team_id?: string | null } & Record }; + teamAlias: string | null; accessToken: string | null; isEditing: boolean; isSaving: boolean; @@ -341,6 +342,7 @@ const ChipList: React.FC<{ values: unknown; emptyLabel: string }> = ({ values, e const ModelInfoEditForm: React.FC = ({ localModelData, modelData, + teamAlias, accessToken, isEditing, isSaving, @@ -799,8 +801,12 @@ const ModelInfoEditForm: React.FC = ({
- Team ID - {modelData.model_info.team_id || "Not Set"} + Team + + {teamAlias + ? `${teamAlias} (${modelData.model_info.team_id})` + : modelData.model_info.team_id || "Not Set"} +
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 3db9418dfb9..f714b8e5c4a 100644 --- a/ui/litellm-dashboard/src/components/model_info_view.test.tsx +++ b/ui/litellm-dashboard/src/components/model_info_view.test.tsx @@ -42,6 +42,11 @@ vi.mock("@/app/(dashboard)/hooks/models/useModelCostMap", () => ({ useModelCostMap: (...args: any[]) => mockUseModelCostMap(...args), })); +const mockUseTeams = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ + useTeams: () => mockUseTeams(), +})); + const mockUsePtuCostAttributionEnabled = vi.fn(); vi.mock("@/app/(dashboard)/hooks/uiSettings/usePtuCostAttributionEnabled", () => ({ usePtuCostAttributionEnabled: () => mockUsePtuCostAttributionEnabled(), @@ -102,6 +107,7 @@ describe("ModelInfoView", () => { }); vi.clearAllMocks(); mockUsePtuCostAttributionEnabled.mockReturnValue(false); + mockUseTeams.mockReturnValue({ data: undefined, isLoading: false, error: null }); mockUseModelsInfo.mockReturnValue({ data: { @@ -1305,6 +1311,78 @@ describe("ModelInfoView", () => { }); }); + describe("team alias", () => { + const teamModel = { + ...defaultModelData, + model_info: { ...defaultModelData.model_info, team_id: "team-1" }, + }; + + beforeEach(() => { + mockUseModelsInfo.mockReturnValue({ data: { data: [teamModel] }, isLoading: false, error: null }); + mockModelInfoV1Call.mockResolvedValue({ data: [teamModel] }); + }); + + const readRawJson = async (user: ReturnType) => { + await user.click(await screen.findByRole("tab", { name: /raw json/i })); + const pre = await screen.findByText(/"model_name": "GPT-4"/, { selector: "pre" }); + return JSON.parse(pre.textContent ?? ""); + }; + + it("shows the team alias next to the team id and adds team_alias to the raw JSON", async () => { + mockUseTeams.mockReturnValue({ + data: [ + { team_id: "team-0", team_alias: "other" }, + { team_id: "team-1", team_alias: "alpha" }, + ], + isLoading: false, + error: null, + }); + const user = userEvent.setup(); + render(, { wrapper }); + + expect(await screen.findByText("alpha (team-1)")).toBeInTheDocument(); + + const raw = await readRawJson(user); + expect(raw.model_info).toMatchObject({ team_id: "team-1", team_alias: "alpha" }); + const keys = Object.keys(raw.model_info); + expect(keys.indexOf("team_alias")).toBe(keys.indexOf("team_id") + 1); + }); + + it("falls back to the bare team id when the team is not in the caller's team list", async () => { + mockUseTeams.mockReturnValue({ + data: [{ team_id: "team-0", team_alias: "other" }], + isLoading: false, + error: null, + }); + const user = userEvent.setup(); + render(, { wrapper }); + + expect(await screen.findByText("team-1")).toBeInTheDocument(); + + const raw = await readRawJson(user); + expect(raw.model_info.team_id).toBe("team-1"); + expect(raw.model_info).not.toHaveProperty("team_alias"); + }); + + it("shows Not Set and no team_alias for a model without a team", async () => { + mockUseModelsInfo.mockReturnValue({ data: { data: [defaultModelData] }, isLoading: false, error: null }); + mockModelInfoV1Call.mockResolvedValue({ data: [defaultModelData] }); + mockUseTeams.mockReturnValue({ + data: [{ team_id: "team-1", team_alias: "alpha" }], + isLoading: false, + error: null, + }); + const user = userEvent.setup(); + render(, { wrapper }); + + expect(await screen.findByText("Team")).toBeInTheDocument(); + expect(screen.queryByText(/alpha/)).not.toBeInTheDocument(); + + const raw = await readRawJson(user); + expect(raw.model_info).not.toHaveProperty("team_alias"); + }); + }); + it("renders the provider card logo from the bundled provider map", async () => { render(, { wrapper }); diff --git a/ui/litellm-dashboard/src/components/model_info_view.tsx b/ui/litellm-dashboard/src/components/model_info_view.tsx index 35afcdb2985..f25416327c0 100644 --- a/ui/litellm-dashboard/src/components/model_info_view.tsx +++ b/ui/litellm-dashboard/src/components/model_info_view.tsx @@ -169,6 +169,12 @@ export default function ModelInfoView({ // Keep modelData variable name for backwards compatibility const modelData = transformedModelData; + const teamAlias = teams?.find((team) => team.team_id === modelData?.model_info?.team_id)?.team_alias || null; + const rawModelInfoEntries = Object.entries(modelData?.model_info ?? {}).flatMap((entry) => + entry[0] === "team_id" && teamAlias ? [entry, ["team_alias", teamAlias]] : [entry], + ); + const rawModelData = modelData && { ...modelData, model_info: Object.fromEntries(rawModelInfoEntries) }; + const canEditModel = canModifyModel({ userRole, userID, isViewOnly }, teams ?? null, { teamId: modelData?.model_info?.team_id, isDbModel: modelData?.model_info?.db_model === true, @@ -765,6 +771,7 @@ export default function ModelInfoView({ - 
{JSON.stringify(modelData, null, 2)}
+ 
+                {JSON.stringify(rawModelData, null, 2)}
+
From d0a846c8be5ec1ae9036254eeb04575f6b406921 Mon Sep 17 00:00:00 2001 From: yassin Date: Sun, 13 Sep 2026 09:27:33 +0000 Subject: [PATCH 027/112] test(router): cover team_model_has_alternatives directly in the mapped router test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/test_router.py | 34 +++++++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index b8a0d70f5bc..7af60e01f7f 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -896,6 +896,40 @@ def test_arouter_test_team_model(): assert result is not None +def test_team_model_has_alternatives(): + def team_deployment(deployment_id: str, team_id: str, public_model_name: str): + return { + "model_name": f"model_name_{team_id}_{deployment_id}", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}, + "model_info": { + "id": deployment_id, + "team_id": team_id, + "team_public_model_name": public_model_name, + }, + } + + router = litellm.Router( + model_list=[ + team_deployment("team-a-1", "team-a", "shared-model"), + team_deployment("team-a-2", "team-a", "shared-model"), + team_deployment("team-a-solo", "team-a", "solo-model"), + team_deployment("team-b-1", "team-b", "shared-model"), + { + "model_name": "plain-model", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}, + "model_info": {"id": "plain-1"}, + }, + ], + ) + + assert router.team_model_has_alternatives("team-a-1") is True + assert router.team_model_has_alternatives("team-a-2") is True + assert router.team_model_has_alternatives("team-a-solo") is False + assert router.team_model_has_alternatives("team-b-1") is False + assert router.team_model_has_alternatives("plain-1") is False + assert router.team_model_has_alternatives("missing-deployment") is False + + def test_arouter_ignore_invalid_deployments(): """ Test that router.ignore_invalid_deployments is set to True From 10f411e60dd7a771a2b3199e3306d197a3127bea Mon Sep 17 00:00:00 2001 From: yassin Date: Sun, 13 Sep 2026 09:34:16 +0000 Subject: [PATCH 028/112] fix(router): name the all-deployments-in-cooldown error on 429 responses RouterRateLimitError now carries the model group's deployment ids so it can tell when every deployment is cooled down, and exposes that as type=all_deployments_in_cooldown with an explicit message. A partial cooldown keeps type=rate_limit_error. Either way the proxy no longer reports type=internal_server_error next to code 429 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 4 ++ litellm/router_utils/handle_error.py | 1 + litellm/types/router.py | 21 ++++++++- .../proxy/test_common_request_processing.py | 33 +++++++++++++ tests/test_litellm/test_router.py | 46 +++++++++++++++++++ 5 files changed, 104 insertions(+), 1 deletion(-) diff --git a/litellm/router.py b/litellm/router.py index 8865543badd..0983c9689c3 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -13878,6 +13878,7 @@ class Router: cooldown_time=_cooldown_time, enable_pre_call_checks=self.enable_pre_call_checks, cooldown_list=_cooldown_list, + model_ids=model_ids, ) if strategy == "simple-shuffle": @@ -13910,6 +13911,7 @@ class Router: cooldown_time=_cooldown_time, enable_pre_call_checks=self.enable_pre_call_checks, cooldown_list=_cooldown_list, + model_ids=model_ids, ) self._override_selector_pre_call_check(strategy, strategy_selector, deployment) verbose_router_logger.info( @@ -14024,6 +14026,7 @@ class Router: cooldown_time=_cooldown_time, enable_pre_call_checks=self.enable_pre_call_checks, cooldown_list=_cooldown_list, + model_ids=model_ids, ) # 6. Apply load balancing strategy @@ -14057,6 +14060,7 @@ class Router: cooldown_time=_cooldown_time, enable_pre_call_checks=self.enable_pre_call_checks, cooldown_list=_cooldown_list, + model_ids=model_ids, ) self._override_selector_pre_call_check(strategy, strategy_selector, deployment) diff --git a/litellm/router_utils/handle_error.py b/litellm/router_utils/handle_error.py index 0e7490d31b1..bfe02675162 100644 --- a/litellm/router_utils/handle_error.py +++ b/litellm/router_utils/handle_error.py @@ -93,4 +93,5 @@ async def async_raise_no_deployment_exception( cooldown_time=_cooldown_time, enable_pre_call_checks=litellm_router_instance.enable_pre_call_checks, cooldown_list=cooldown_list_ids, + model_ids=model_ids, ) diff --git a/litellm/types/router.py b/litellm/types/router.py index c7363502017..ddc13e18567 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -645,6 +645,7 @@ class RouterErrors(enum.Enum): user_defined_ratelimit_error = "Deployment over user-defined ratelimit." no_deployments_available = "No deployments available for selected model" + all_deployments_in_cooldown = "All deployments for selected model are in cooldown" no_deployments_with_tag_routing = "Not allowed to access model due to tags configuration" no_deployments_with_provider_budget_routing = "No deployments available - crossed budget" no_healthy_deployments = "There are no healthy deployments for this model" @@ -868,6 +869,11 @@ class RouterRateLimitErrorBasic(ValueError): super().__init__(_message) +class RouterErrorTypes(str, enum.Enum): + rate_limit_error = "rate_limit_error" + all_deployments_in_cooldown = "all_deployments_in_cooldown" + + class RouterRateLimitError(ValueError): def __init__( self, @@ -875,12 +881,25 @@ class RouterRateLimitError(ValueError): cooldown_time: float, enable_pre_call_checks: bool, cooldown_list: list, + model_ids: Sequence[str] = (), ) -> None: self.model = model self.cooldown_time = cooldown_time self.enable_pre_call_checks = enable_pre_call_checks self.cooldown_list = cooldown_list - _message = f"{RouterErrors.no_deployments_available.value}, Try again in {cooldown_time} seconds. Passed model={model}. pre-call-checks={enable_pre_call_checks}, cooldown_list={cooldown_list}" + self.all_deployments_in_cooldown = bool(model_ids) and frozenset(model_ids) <= frozenset(cooldown_list) + self.type = ( + RouterErrorTypes.all_deployments_in_cooldown.value + if self.all_deployments_in_cooldown + else RouterErrorTypes.rate_limit_error.value + ) + _reason: Final = ( + f" {RouterErrors.all_deployments_in_cooldown.value}." if self.all_deployments_in_cooldown else "" + ) + _message: Final = ( + f"{RouterErrors.no_deployments_available.value}, Try again in {cooldown_time} seconds.{_reason} " + f"Passed model={model}. pre-call-checks={enable_pre_call_checks}, cooldown_list={cooldown_list}" + ) super().__init__(_message) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 69e89d1c604..92b80685a1f 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -3879,6 +3879,39 @@ class TestHandleLLMApiExceptionRetryAfter: assert proxy_exc.headers["retry-after"] == "43" assert proxy_exc.headers["x-custom"] == "1" + async def test_handle_llm_api_exception_names_cooldown_when_every_deployment_is_cooled_down(self): + from litellm.types.router import RouterRateLimitError + + exc = RouterRateLimitError( + model="gpt-4", + cooldown_time=120, + enable_pre_call_checks=False, + cooldown_list=["dep-a", "dep-b"], + model_ids=["dep-a", "dep-b"], + ) + proxy_exc = await self._invoke(exc) + body = proxy_exc.to_dict() + assert body["type"] == "all_deployments_in_cooldown" + assert body["code"] == "429" + assert "All deployments for selected model are in cooldown" in body["message"] + assert proxy_exc.headers["retry-after"] == "120" + + async def test_handle_llm_api_exception_keeps_rate_limit_type_when_cooldown_is_partial(self): + from litellm.types.router import RouterRateLimitError + + exc = RouterRateLimitError( + model="gpt-4", + cooldown_time=120, + enable_pre_call_checks=False, + cooldown_list=["dep-a"], + model_ids=["dep-a", "dep-b"], + ) + proxy_exc = await self._invoke(exc) + body = proxy_exc.to_dict() + assert body["type"] == "rate_limit_error" + assert body["code"] == "429" + assert "All deployments for selected model are in cooldown" not in body["message"] + class TestHandleLLMApiExceptionFramingHeaders: """HTTP-framing headers on the provider exception must be stripped before the diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index b8a0d70f5bc..cc571ad3d3c 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -7735,6 +7735,52 @@ def test_get_available_deployment_raises_when_addressed_dict_is_blocked(): router.get_available_deployment(model="dep-0", request_kwargs={}) +def _cool_down(router: Router, *deployment_ids: str) -> None: + for deployment_id in deployment_ids: + router.cooldown_cache.add_deployment_to_cooldown( + model_id=deployment_id, + original_exception=litellm.RateLimitError(message="upstream 429", llm_provider="openai", model="gpt-4o"), + exception_status=429, + cooldown_time=60, + ) + + +async def _select_deployment(router: Router, use_async: bool) -> None: + if use_async: + await router.async_get_available_deployment(model="gpt-4o", request_kwargs={}) + return + router.get_available_deployment(model="gpt-4o", request_kwargs={}) + + +@pytest.mark.parametrize("use_async", [False, True], ids=["sync", "async"]) +@pytest.mark.asyncio +async def test_get_available_deployment_names_cooldown_when_every_deployment_is_cooled_down(use_async: bool): + from litellm.types.router import RouterErrors, RouterRateLimitError + + router: Final = _router_with_two_deployments([False, False]) + _cool_down(router, "dep-0", "dep-1") + with pytest.raises(RouterRateLimitError) as exc_info: + await _select_deployment(router, use_async) + assert exc_info.value.all_deployments_in_cooldown is True + assert exc_info.value.type == "all_deployments_in_cooldown" + assert RouterErrors.all_deployments_in_cooldown.value in str(exc_info.value) + assert str(exc_info.value).startswith("No deployments available for selected model, Try again in ") + + +@pytest.mark.parametrize("use_async", [False, True], ids=["sync", "async"]) +@pytest.mark.asyncio +async def test_get_available_deployment_keeps_generic_error_when_cooldown_is_partial(use_async: bool): + from litellm.types.router import RouterErrors, RouterRateLimitError + + router: Final = _router_with_two_deployments([False, True]) + _cool_down(router, "dep-0") + with pytest.raises(RouterRateLimitError) as exc_info: + await _select_deployment(router, use_async) + assert exc_info.value.all_deployments_in_cooldown is False + assert exc_info.value.type == "rate_limit_error" + assert RouterErrors.all_deployments_in_cooldown.value not in str(exc_info.value) + + def _router_with_two_pass_through_deployments(blocked_flags): import litellm From 9080f0904ad6bee6c5f10debdccdd1b064e1efe3 Mon Sep 17 00:00:00 2001 From: yassin Date: Sun, 13 Sep 2026 09:35:54 +0000 Subject: [PATCH 029/112] fix(router): ignore blocked siblings when checking team model cooldown alternatives Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 4 +++- .../router_utils/test_cooldown_handlers.py | 18 +++++++++++++++++- tests/test_litellm/test_router.py | 6 +++++- 3 files changed, 25 insertions(+), 3 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 4929b17f7fc..63d5f86d46a 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1561,7 +1561,9 @@ class Router: public_model_name: Final = deployment.model_info.team_public_model_name if team_id is None or public_model_name is None: return False - return len(self.team_model_to_deployment_indices.get((team_id, public_model_name)) or ()) > 1 + sibling_indices: Final = self.team_model_to_deployment_indices.get((team_id, public_model_name)) or () + routable_siblings: Final = self._filter_blocked_deployments([self.model_list[idx] for idx in sibling_indices]) + return len(routable_siblings) > 1 _OVERRIDABLE_ROUTING_STRATEGIES: frozenset[str] = frozenset({"simple-shuffle", *_DEFAULT_SELECTOR_ATTR_BY_STRATEGY}) diff --git a/tests/test_litellm/router_utils/test_cooldown_handlers.py b/tests/test_litellm/router_utils/test_cooldown_handlers.py index 6f66bb863cc..fdbfd618dab 100644 --- a/tests/test_litellm/router_utils/test_cooldown_handlers.py +++ b/tests/test_litellm/router_utils/test_cooldown_handlers.py @@ -440,7 +440,7 @@ class TestRoutingGroupCooldownAlternatives: class TestTeamModelCooldownAlternatives: - def _router(self, team_deployments: int): + def _router(self, team_deployments: int, blocked_ids: frozenset[str] = frozenset()): from litellm import Router return Router( @@ -452,6 +452,7 @@ class TestTeamModelCooldownAlternatives: "id": f"team-deploy-{i}", "team_id": "team-1", "team_public_model_name": "team-gpt-4o-mini", + "blocked": f"team-deploy-{i}" in blocked_ids, }, } for i in range(team_deployments) @@ -487,3 +488,18 @@ class TestTeamModelCooldownAlternatives: ) is False ) + + def test_429_with_only_a_blocked_sibling_keeps_single_deployment_exemption(self): + from litellm.router_utils.cooldown_handlers import _should_cooldown_deployment + + router = self._router(team_deployments=2, blocked_ids=frozenset({"team-deploy-1"})) + assert ( + _should_cooldown_deployment( + litellm_router_instance=router, + deployment="team-deploy-0", + exception_status=429, + original_exception=Exception("rate limited"), + requested_model_group="team-gpt-4o-mini", + ) + is False + ) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 7af60e01f7f..7786bfb4788 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -897,7 +897,7 @@ def test_arouter_test_team_model(): def test_team_model_has_alternatives(): - def team_deployment(deployment_id: str, team_id: str, public_model_name: str): + def team_deployment(deployment_id: str, team_id: str, public_model_name: str, blocked: bool = False): return { "model_name": f"model_name_{team_id}_{deployment_id}", "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}, @@ -905,6 +905,7 @@ def test_team_model_has_alternatives(): "id": deployment_id, "team_id": team_id, "team_public_model_name": public_model_name, + "blocked": blocked, }, } @@ -914,6 +915,8 @@ def test_team_model_has_alternatives(): team_deployment("team-a-2", "team-a", "shared-model"), team_deployment("team-a-solo", "team-a", "solo-model"), team_deployment("team-b-1", "team-b", "shared-model"), + team_deployment("team-c-1", "team-c", "paused-sibling-model"), + team_deployment("team-c-paused", "team-c", "paused-sibling-model", blocked=True), { "model_name": "plain-model", "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}, @@ -926,6 +929,7 @@ def test_team_model_has_alternatives(): assert router.team_model_has_alternatives("team-a-2") is True assert router.team_model_has_alternatives("team-a-solo") is False assert router.team_model_has_alternatives("team-b-1") is False + assert router.team_model_has_alternatives("team-c-1") is False assert router.team_model_has_alternatives("plain-1") is False assert router.team_model_has_alternatives("missing-deployment") is False From e61b6bfd5ff58c61381e32aae075b2516fd3c76d Mon Sep 17 00:00:00 2001 From: yassin Date: Sun, 13 Sep 2026 09:48:35 +0000 Subject: [PATCH 030/112] fix(router): classify pass-through cooldown against pass-through deployments only Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 7 ++++++- tests/test_litellm/test_router.py | 18 ++++++++++++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/litellm/router.py b/litellm/router.py index 0983c9689c3..e74b2079fd4 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -13989,6 +13989,11 @@ class Router: model=model, llm_provider="", ) + pass_through_model_ids: Final = tuple( + deployment["model_info"]["id"] + for deployment in pass_through_deployments + if "id" in deployment.get("model_info", {}) + ) # 4. Apply health-check and cooldown filtering parent_otel_span: Final[Span | None] = _get_parent_otel_span_from_kwargs(request_kwargs) @@ -14026,7 +14031,7 @@ class Router: cooldown_time=_cooldown_time, enable_pre_call_checks=self.enable_pre_call_checks, cooldown_list=_cooldown_list, - model_ids=model_ids, + model_ids=pass_through_model_ids, ) # 6. Apply load balancing strategy diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index cc571ad3d3c..35d0b8104bd 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -7818,6 +7818,24 @@ def test_get_available_deployment_for_pass_through_raises_when_dict_blocked(): ) +def test_get_available_deployment_for_pass_through_names_cooldown_despite_healthy_non_pass_through(): + from litellm.types.router import RouterRateLimitError + + router: Final = _router_with_two_pass_through_deployments([False, False]) + router.add_deployment( + Deployment( + model_name="gpt-4o", + litellm_params=LiteLLM_Params(model="openai/gpt-4o-plain", api_key="sk-fake-for-tests"), + model_info=ModelInfo(id="plain-0"), + ) + ) + _cool_down(router, "pt-0", "pt-1") + with pytest.raises(RouterRateLimitError) as exc_info: + router.get_available_deployment_for_pass_through(model="gpt-4o", request_kwargs={}) + assert exc_info.value.all_deployments_in_cooldown is True + assert exc_info.value.type == "all_deployments_in_cooldown" + + def test_initialize_deployment_for_pass_through_keeps_bedrock_iam_deployment(): """ Bedrock deployments using IAM/OIDC auth have no api_key; pass-through From e41b3bd13fa6415de7b4076dda82f673fac8b957 Mon Sep 17 00:00:00 2001 From: yassin Date: Sun, 13 Sep 2026 09:50:12 +0000 Subject: [PATCH 031/112] test(router): annotate return types of team cooldown test helpers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/router_utils/test_cooldown_handlers.py | 6 ++---- tests/test_litellm/test_router.py | 4 +++- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/test_litellm/router_utils/test_cooldown_handlers.py b/tests/test_litellm/router_utils/test_cooldown_handlers.py index fdbfd618dab..6fed4be5909 100644 --- a/tests/test_litellm/router_utils/test_cooldown_handlers.py +++ b/tests/test_litellm/router_utils/test_cooldown_handlers.py @@ -440,10 +440,8 @@ class TestRoutingGroupCooldownAlternatives: class TestTeamModelCooldownAlternatives: - def _router(self, team_deployments: int, blocked_ids: frozenset[str] = frozenset()): - from litellm import Router - - return Router( + def _router(self, team_deployments: int, blocked_ids: frozenset[str] = frozenset()) -> litellm.Router: + return litellm.Router( model_list=[ { "model_name": f"model_name_team-1_{i}", diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 7786bfb4788..c6b3b7fb8d6 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -897,7 +897,9 @@ def test_arouter_test_team_model(): def test_team_model_has_alternatives(): - def team_deployment(deployment_id: str, team_id: str, public_model_name: str, blocked: bool = False): + def team_deployment( + deployment_id: str, team_id: str, public_model_name: str, blocked: bool = False + ) -> DeploymentTypedDict: return { "model_name": f"model_name_{team_id}_{deployment_id}", "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}, From db79226b6b8786b40d10e8736595a0fe6bf07f47 Mon Sep 17 00:00:00 2001 From: mateo Date: Sun, 13 Sep 2026 09:59:39 +0000 Subject: [PATCH 032/112] test(auth): freeze the cache clock in auth prefetch tests The org cache entries written by prefetch_auth_objects carry the 5s DEFAULT_IN_MEMORY_TTL. The first @log_db_metrics getter lazily imports litellm.proxy.proxy_server, which on a cold CI runner can take longer than 5s, so the org entry expired before get_org_object read it and the getter fell through to the MagicMock database. Inject a frozen clock into InMemoryCache so the test asserts the join, not import latency. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy_behavior/auth/test_auth_object_prefetch.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/tests/proxy_behavior/auth/test_auth_object_prefetch.py b/tests/proxy_behavior/auth/test_auth_object_prefetch.py index e2d947f4284..59e9585a296 100644 --- a/tests/proxy_behavior/auth/test_auth_object_prefetch.py +++ b/tests/proxy_behavior/auth/test_auth_object_prefetch.py @@ -22,6 +22,11 @@ from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache pytestmark = pytest.mark.asyncio(loop_scope="session") +def _frozen_cache() -> UserApiKeyCache: + """The org entries carry a 5s TTL; a frozen clock keeps a slow first call from expiring them mid-test.""" + return UserApiKeyCache(in_memory_cache=InMemoryCache(clock=lambda: 1_000_000.0), redis_cache=None) + + def _dead_db() -> MagicMock: prisma = MagicMock(name="prisma_client") prisma.db.query_first = AsyncMock(return_value=None) @@ -58,7 +63,7 @@ async def test_join_binds_the_membership_to_the_requested_team(prisma): data={"user_id": user_id, "team_id": team_b, "litellm_budget_table": {"connect": {"budget_id": f"b-{run}"}}} ) - cache = UserApiKeyCache(in_memory_cache=InMemoryCache(), redis_cache=None) + cache = _frozen_cache() refs = AuthObjectRefs(user_id=user_id, team_id=team_a, membership_user_id=user_id, organization_id=org_id) await prefetch_auth_objects(refs=refs, user_api_key_cache=cache, prisma_client=prisma) @@ -100,7 +105,7 @@ async def test_join_reads_team_model_aliases_from_the_mapped_column(prisma): where={"team_id": team_id}, include={"litellm_model_table": True} ) - cache = UserApiKeyCache(in_memory_cache=InMemoryCache(), redis_cache=None) + cache = _frozen_cache() refs = AuthObjectRefs(user_id=None, team_id=team_id, membership_user_id=None, organization_id=None) await prefetch_auth_objects(refs=refs, user_api_key_cache=cache, prisma_client=prisma) @@ -144,7 +149,7 @@ async def test_join_reads_null_nested_lists_the_way_prisma_does(prisma): where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, include={"litellm_budget_table": True} ) - cache = UserApiKeyCache(in_memory_cache=InMemoryCache(), redis_cache=None) + cache = _frozen_cache() refs = AuthObjectRefs(user_id=user_id, team_id=team_id, membership_user_id=user_id, organization_id=None) await prefetch_auth_objects(refs=refs, user_api_key_cache=cache, prisma_client=prisma) From a70f9a17ccb2e4dae4bb7a90236d514a45bc9e5b Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 14 Sep 2026 08:54:34 +0000 Subject: [PATCH 033/112] perf(logging): skip correlation contextvar stamping when request_correlation_in_logs is off Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/litellm_logging.py | 8 +- .../test_litellm_logging.py | 90 +++++++++++++++++-- .../test_streaming_handler.py | 24 +++-- 3 files changed, 104 insertions(+), 18 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 5fe94320491..9ba9fd082f3 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -556,7 +556,7 @@ class Logging(LiteLLMLoggingBaseClass): # ids leaking into a different, later request on the same thread. Sync # support is deferred to a follow-up PR with its own safe-restore # mechanism; async calls (the proxy's only call path) are unaffected. - if supports_correlation_logging: + if supports_correlation_logging and litellm.request_correlation_in_logs: set_trace_id(self.litellm_trace_id) set_session_id(self.litellm_session_id) # set_trace_id()/set_session_id() sanitize (strip control chars, bound @@ -2442,7 +2442,7 @@ class Logging(LiteLLMLoggingBaseClass): call) would leave the outer request's subsequent log lines stamped with the nested call's trace_id/session_id instead of its own. - Uses a plain set() of the captured pre-call value rather than + Uses a plain contextvar set() of the captured pre-call value rather than contextvars.Token-based reset(), since this can end up called from a different asyncio Task/context than __init__ ran in (e.g. the request task's own wrapper() finally block, plus async_success_handler @@ -2453,8 +2453,8 @@ class Logging(LiteLLMLoggingBaseClass): that Task's view of the contextvars, so calling it multiple times (once per Task involved in this attempt) is required, not just safe. """ - set_trace_id(self._pre_call_trace_id) - set_session_id(self._pre_call_session_id) + trace_id_var.set(self._pre_call_trace_id) + session_id_var.set(self._pre_call_session_id) def _restore_correlation_context_if_unclaimed(self) -> None: """Guarded variant for __del__-triggered cleanup only. diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 451b740dd71..70f9bae283b 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -5221,10 +5221,11 @@ def test_handle_anthropic_messages_parsed_response_logging_preserves_fast_mode_s assert getattr(result.usage, "speed", None) == "fast" -def test_logging_init_sets_trace_id(): +def test_logging_init_sets_trace_id(monkeypatch): """Logging.__init__() must call set_trace_id with self.litellm_trace_id.""" from litellm.litellm_core_utils.litellm_logging import Logging + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) trace_id_var.set("") log_obj = Logging( @@ -5240,7 +5241,7 @@ def test_logging_init_sets_trace_id(): assert trace_id_var.get() == log_obj.litellm_trace_id -def test_logging_init_skips_stamping_when_correlation_logging_unsupported(): +def test_logging_init_skips_stamping_when_correlation_logging_unsupported(monkeypatch): """supports_correlation_logging=False (what wrapper(), the sync entry point, always passes) must leave trace_id_var/session_id_var completely untouched, even though self.litellm_trace_id/litellm_session_id (the @@ -5248,6 +5249,7 @@ def test_logging_init_skips_stamping_when_correlation_logging_unsupported(): usual - only the ambient contextvar stamping is gated.""" from litellm.litellm_core_utils.litellm_logging import Logging + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) trace_id_var.set("") session_id_var.set("") @@ -5271,10 +5273,48 @@ def test_logging_init_skips_stamping_when_correlation_logging_unsupported(): assert log_obj.litellm_session_id == "should-not-be-stamped" -def test_logging_init_sets_session_id_when_provided(): +def test_logging_init_skips_stamping_when_request_correlation_in_logs_disabled(monkeypatch): + from litellm.litellm_core_utils.litellm_logging import Logging + + monkeypatch.setattr(litellm, "request_correlation_in_logs", False) + trace_id_var.set("outer") + session_id_var.set("outer-sid") + try: + with ( + patch( # test-quality-ok: regression test verifies disabled stamping skips both setters + "litellm.litellm_core_utils.litellm_logging.set_trace_id" + ) as mock_set_trace_id, + patch( # test-quality-ok: regression test verifies disabled stamping skips both setters + "litellm.litellm_core_utils.litellm_logging.set_session_id" + ) as mock_set_session_id, + ): + log_obj = Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="completion", + start_time=None, + litellm_call_id="call-disabled", + function_id="fn-disabled", + kwargs={"litellm_session_id": "disabled-session"}, + supports_correlation_logging=True, + ) + + assert trace_id_var.get() == "outer" + assert session_id_var.get() == "outer-sid" + assert log_obj._own_trace_id == "outer" + mock_set_trace_id.assert_not_called() + mock_set_session_id.assert_not_called() + finally: + trace_id_var.set("") + session_id_var.set("") + + +def test_logging_init_sets_session_id_when_provided(monkeypatch): """Logging.__init__() must call set_session_id when litellm_session_id is in kwargs.""" from litellm.litellm_core_utils.litellm_logging import Logging + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) session_id_var.set("") Logging( @@ -5290,11 +5330,12 @@ def test_logging_init_sets_session_id_when_provided(): assert session_id_var.get() == "my-session-99" -def test_logging_init_resets_session_id_to_empty_when_absent(): +def test_logging_init_resets_session_id_to_empty_when_absent(monkeypatch): """When no session_id is in kwargs, Logging.__init__() must reset session_id_var to "" so a prior request's session_id does not leak into subsequent log records.""" from litellm.litellm_core_utils.litellm_logging import Logging + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) session_id_var.set("preexisting-sid") Logging( @@ -5310,7 +5351,7 @@ def test_logging_init_resets_session_id_to_empty_when_absent(): assert session_id_var.get() == "" -def test_restore_correlation_context_resets_to_pre_call_value(): +def test_restore_correlation_context_resets_to_pre_call_value(monkeypatch): """_restore_correlation_context() must put trace_id_var/session_id_var back to whatever they were immediately before this Logging instance was constructed. This is the mechanism that prevents a nested call (e.g. a guardrail's own @@ -5318,6 +5359,7 @@ def test_restore_correlation_context_resets_to_pre_call_value(): session_id into the outer call's subsequent log lines.""" from litellm.litellm_core_utils.litellm_logging import Logging + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) trace_id_var.set("outer-trace") session_id_var.set("outer-session") try: @@ -5343,7 +5385,7 @@ def test_restore_correlation_context_resets_to_pre_call_value(): session_id_var.set("") -def test_restore_correlation_context_safe_to_call_repeatedly(): +def test_restore_correlation_context_safe_to_call_repeatedly(monkeypatch): """Calling _restore_correlation_context() more than once must not raise. It's deliberately NOT guarded against repeat calls: wrapper()'s finally @@ -5353,6 +5395,7 @@ def test_restore_correlation_context_safe_to_call_repeatedly(): the contextvars, so repeat calls are expected, not just tolerated.""" from litellm.litellm_core_utils.litellm_logging import Logging + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) log_obj = Logging( model="gpt-3.5-turbo", messages=[{"role": "user", "content": "hi"}], @@ -5367,8 +5410,40 @@ def test_restore_correlation_context_safe_to_call_repeatedly(): log_obj._restore_correlation_context() # must not raise +def test_restore_correlation_context_does_not_resanitize(monkeypatch): + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm._logging import _sanitize_correlation_id + + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + trace_id_var.set("outer-trace") + session_id_var.set("outer-session") + try: + log_obj = Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="completion", + start_time=None, + litellm_call_id="call-no-resanitize", + function_id="fn-no-resanitize", + kwargs={"litellm_session_id": "inner-session"}, + ) + + with patch( # test-quality-ok: regression test verifies restore avoids sanitization + "litellm._logging._sanitize_correlation_id", wraps=_sanitize_correlation_id + ) as mock_sanitize: + log_obj._restore_correlation_context() + + mock_sanitize.assert_not_called() + assert trace_id_var.get() == "outer-trace" + assert session_id_var.get() == "outer-session" + finally: + trace_id_var.set("") + session_id_var.set("") + + @pytest.mark.asyncio -async def test_restore_correlation_context_works_across_asyncio_task_boundary(): +async def test_restore_correlation_context_works_across_asyncio_task_boundary(monkeypatch): """_restore_correlation_context() must succeed even when it's called from a different asyncio Task than the one Logging.__init__() ran in - exactly what happens on litellm's real async success path, where async_success_handler is @@ -5385,6 +5460,7 @@ async def test_restore_correlation_context_works_across_asyncio_task_boundary(): """ from litellm.litellm_core_utils.litellm_logging import Logging + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) trace_id_var.set("outer-trace-cross-task") session_id_var.set("outer-session-cross-task") try: diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 37e2031fdf4..784468c839b 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -4100,7 +4100,7 @@ async def test_async_streaming_completion_does_not_reset_context_before_iteratio session_id_var.set("") -def test_stream_wrapper_del_restores_correlation_context(): +def test_stream_wrapper_del_restores_correlation_context(monkeypatch): """CustomStreamWrapper.__del__ is the best-effort fallback for an abandoned stream (caller never exhausts it, so the normal terminal-handler restore never fires). Testing this via real garbage collection is unreliable in @@ -4112,6 +4112,7 @@ def test_stream_wrapper_del_restores_correlation_context(): doesn't run actual finalization, and this exercises exactly the logic that real garbage collection would eventually trigger. """ + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) trace_id_var.set("outer-trace-abandoned") session_id_var.set("outer-session-abandoned") try: @@ -4159,12 +4160,13 @@ def test_stream_wrapper_del_never_raises_with_broken_logging_obj(): wrapper.__del__() # must not raise -def test_stream_wrapper_del_does_not_clobber_a_newer_active_call(): +def test_stream_wrapper_del_does_not_clobber_a_newer_active_call(monkeypatch): """A delayed finalizer must never stomp a different, still-active call's context. If an abandoned stream's __del__ fires late - after a new call has already started in the same Task/thread and claimed the contextvars - unconditionally restoring the abandoned stream's own pre-call snapshot would corrupt the active call's subsequent log lines with stale ids.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) trace_id_var.set("outer-trace-before-abandoned-call") session_id_var.set("outer-session-before-abandoned-call") try: @@ -4210,13 +4212,14 @@ def test_stream_wrapper_del_does_not_clobber_a_newer_active_call(): session_id_var.set("") -def test_stream_wrapper_del_restores_when_own_session_id_needed_sanitizing(): +def test_stream_wrapper_del_restores_when_own_session_id_needed_sanitizing(monkeypatch): """The __del__ guard must compare against the *sanitized* id actually stored in the contextvar, not the raw litellm_session_id/litellm_trace_id - set_session_id()/set_trace_id() strip control characters before storing, so a caller-supplied id containing e.g. a newline would never equal the raw attribute, and the guard would wrongly conclude some other call has claimed the context and skip cleanup forever.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) trace_id_var.set("outer-trace-needs-sanitizing") session_id_var.set("outer-session-needs-sanitizing") try: @@ -4250,7 +4253,7 @@ def test_stream_wrapper_del_restores_when_own_session_id_needed_sanitizing(): session_id_var.set("") -def test_stream_wrapper_next_keeps_context_active_through_synthesized_finish_reason_chunk(): +def test_stream_wrapper_next_keeps_context_active_through_synthesized_finish_reason_chunk(monkeypatch): """When the underlying stream ends without ever emitting an explicit finish_reason chunk, __next__ synthesizes one via finish_reason_handler() and returns it. That chunk is still this call's own data - the caller's @@ -4261,6 +4264,7 @@ def test_stream_wrapper_next_keeps_context_active_through_synthesized_finish_rea correct, deterministic restore on the very next __next__() call, since completion_stream is already exhausted and immediately re-raises StopIteration.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) trace_id_var.set("outer-trace-finish-reason") session_id_var.set("outer-session-finish-reason") try: @@ -4300,12 +4304,13 @@ def test_stream_wrapper_next_keeps_context_active_through_synthesized_finish_rea session_id_var.set("") -def test_stream_wrapper_del_cleans_up_after_synthesized_finish_reason_chunk(): +def test_stream_wrapper_del_cleans_up_after_synthesized_finish_reason_chunk(monkeypatch): """A caller that breaks immediately after seeing finish_reason (the early-break pattern) never triggers the next()-driven restore above - it relies on the best-effort __del__ guard instead, same as any other abandoned stream. The guard must still recognize this call's own (unrestored) ids as unclaimed and clean them up.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) trace_id_var.set("outer-trace-finish-reason-del") session_id_var.set("outer-session-finish-reason-del") try: @@ -4338,10 +4343,11 @@ def test_stream_wrapper_del_cleans_up_after_synthesized_finish_reason_chunk(): @pytest.mark.asyncio -async def test_stream_wrapper_anext_keeps_context_active_through_synthesized_finish_reason_chunk(): +async def test_stream_wrapper_anext_keeps_context_active_through_synthesized_finish_reason_chunk(monkeypatch): """Async sibling of test_stream_wrapper_next_keeps_context_active_through_synthesized_finish_reason_chunk - _finalize_completed_stream()'s else branch must not restore before returning the synthesized chunk either.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) trace_id_var.set("outer-trace-anext-finish-reason") session_id_var.set("outer-session-anext-finish-reason") try: @@ -4394,6 +4400,7 @@ async def test_stream_wrapper_anext_max_duration_timeout_restores_consumer_corre path as every other failure so the consumer's outer correlation context gets restored - calling the check before entering __anext__()'s try block would let the Timeout bypass that restoration entirely.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) monkeypatch.setattr(litellm.constants, "LITELLM_MAX_STREAMING_DURATION_SECONDS", 1) trace_id_var.set("outer-trace-max-duration") session_id_var.set("outer-session-max-duration") @@ -4434,12 +4441,13 @@ async def test_stream_wrapper_anext_max_duration_timeout_restores_consumer_corre @pytest.mark.asyncio -async def test_stream_wrapper_aclose_restores_consumer_correlation_context(): +async def test_stream_wrapper_aclose_restores_consumer_correlation_context(monkeypatch): """Explicit early termination (aclose(), e.g. on client disconnect or a router fallback aborting an in-progress stream) must restore the caller's correlation context too - not just __del__'s best-effort GC-timed fallback, since aclose() is normally called deterministically by the consumer/ framework, unlike __del__.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) trace_id_var.set("outer-trace-aclose") session_id_var.set("outer-session-aclose") try: @@ -4481,6 +4489,7 @@ async def test_stream_wrapper_aclose_keeps_context_active_through_close_failure_ branch logs a debug diagnostic. That log line must still carry the closing stream's own trace_id/session_id - the outer context must not be restored until after the close attempt (and its diagnostic) completes.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) trace_id_var.set("outer-trace-close-fail") session_id_var.set("outer-session-close-fail") try: @@ -4541,6 +4550,7 @@ def test_handle_stream_fallback_error_restores_context_only_after_exception_mapp mapping. The consumer's outer context must not be restored until that mapping call returns, or the diagnostic log line would carry the outer (or empty) trace_id/session_id instead of the failing stream's own.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) trace_id_var.set("outer-trace-fallback") session_id_var.set("outer-session-fallback") try: From 1661e72c2c53a63c010685c04354bcaabcd12a97 Mon Sep 17 00:00:00 2001 From: mateo Date: Mon, 14 Sep 2026 13:21:46 +0000 Subject: [PATCH 034/112] fix(registry): drop retired friendliai llama-3.1 serverless models Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 28 ------------------- model_prices_and_context_window.json | 28 ------------------- 2 files changed, 56 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 86605d77aab..fa67d8de443 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -23164,34 +23164,6 @@ "output_cost_per_token": 0.0, "source": "https://fireworks.ai/pricing" }, - "friendliai/meta-llama-3.1-70b-instruct": { - "input_cost_per_token": 6e-07, - "litellm_provider": "friendliai", - "max_input_tokens": 8192, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 6e-07, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, - "friendliai/meta-llama-3.1-8b-instruct": { - "input_cost_per_token": 1e-07, - "litellm_provider": "friendliai", - "max_input_tokens": 8192, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1e-07, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, "friendliai/zai-org/GLM-5.3-Flash": { "litellm_provider": "friendliai", "max_input_tokens": 1048576, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 86605d77aab..fa67d8de443 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -23164,34 +23164,6 @@ "output_cost_per_token": 0.0, "source": "https://fireworks.ai/pricing" }, - "friendliai/meta-llama-3.1-70b-instruct": { - "input_cost_per_token": 6e-07, - "litellm_provider": "friendliai", - "max_input_tokens": 8192, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 6e-07, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, - "friendliai/meta-llama-3.1-8b-instruct": { - "input_cost_per_token": 1e-07, - "litellm_provider": "friendliai", - "max_input_tokens": 8192, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1e-07, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, "friendliai/zai-org/GLM-5.3-Flash": { "litellm_provider": "friendliai", "max_input_tokens": 1048576, From 58c3ccf6684ec4cfed0aab4ba6d45e312a9bcce1 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 15:27:06 +0000 Subject: [PATCH 035/112] fix(proxy): keep org admins' own team memberships in other orgs visible on team list An org admin listing their own teams on GET /team/list and GET /v2/team/list only saw teams in the orgs they administer. Teams they belong to in other orgs were dropped because the org scope and the membership scope were ANDed. A self query now unions the two, while a query for another user keeps the org boundary intersection. Resolves LIT-3723 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints/team_endpoints.py | 143 ++++++++++------ .../test_team_endpoints.py | 161 +++++++++++++++++- 2 files changed, 247 insertions(+), 57 deletions(-) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 0b7f69bbb7f..b4608ac27a0 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -179,6 +179,7 @@ from litellm.types.proxy.management_endpoints.team_endpoints import ( if TYPE_CHECKING: from prisma import Prisma from prisma import models as prisma_models + from prisma import types as prisma_types router: Final = APIRouter() @@ -4857,6 +4858,25 @@ async def _get_org_admin_org_ids( return org_ids if org_ids else None +async def _get_user_team_ids( + user_id: str, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> tuple[str, ...]: + try: + user: Final = await get_user_object( + user_id=user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + proxy_logging_obj=proxy_logging_obj, + ) + except ValueError: + return () + return tuple(user.teams or ()) if user is not None else () + + async def _build_team_list_where_conditions( prisma_client: PrismaClient, team_id: str | None, @@ -4867,12 +4887,16 @@ async def _build_team_list_where_conditions( search: str | None = None, search_team_id_match: TeamIdSearchMatch = "exact", org_admin_org_ids: list[str] | None = None, + own_team_ids: tuple[str, ...] = (), user_api_key_cache: UserApiKeyCache | None = None, proxy_logging_obj: ProxyLogging | None = None, ) -> dict[str, object] | None: """ Build where conditions for team list query. + An org admin listing their own teams sees the union of the teams in the + orgs they administer and `own_team_ids`, the teams they are a member of. + Returns None when the query is guaranteed to yield no results (e.g. user has no team memberships), allowing the caller to skip the DB round-trip. """ @@ -4895,6 +4919,11 @@ async def _build_team_list_where_conditions( if organization_id: where_conditions["organization_id"] = organization_id + elif org_admin_org_ids is not None and own_team_ids: + org_or_membership_scope: Final[prisma_types.LiteLLM_TeamTableWhereInput] = { + "OR": [{"organization_id": {"in": org_admin_org_ids}}, {"team_id": {"in": list(own_team_ids)}}] + } + where_conditions["AND"] = [org_or_membership_scope] elif org_admin_org_ids is not None: # Org admin: always scope to their orgs, even when filtering by user_id. where_conditions["organization_id"] = {"in": org_admin_org_ids} @@ -5026,66 +5055,72 @@ async def _enforce_list_team_v2_access( prisma_client: PrismaClient, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, -) -> tuple[str | None, list[str] | None]: +) -> tuple[str | None, list[str] | None, tuple[str, ...]]: """Enforce access control for list_team_v2. - Proxy admins and admin viewers can query any teams. - - Org admins can query teams within their organizations. + - Org admins can query teams within their organizations, plus the teams + they are a member of when listing their own teams. - Regular users can only query their own teams. - Returns the (possibly overridden) user_id and org_admin_org_ids. + Returns the (possibly overridden) user_id, org_admin_org_ids and, for an + org admin's own query, the caller's own team ids. """ is_proxy_admin: Final = _user_has_admin_view(user_api_key_dict) - org_admin_org_ids: list[str] | None = None + caller_user_id: Final = user_api_key_dict.user_id if is_proxy_admin: - return user_id, org_admin_org_ids + return user_id, None, () # Always check org admin status so that even own-queries see # the full set of organisation teams, not just direct memberships. - if user_api_key_dict.user_id: - org_admin_org_ids = await _get_org_admin_org_ids( - user_id=user_api_key_dict.user_id, + org_admin_org_ids: Final = ( + await _get_org_admin_org_ids( + user_id=caller_user_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) + if caller_user_id + else None + ) - if org_admin_org_ids is not None: + if caller_user_id and org_admin_org_ids is not None: # Org admin: validate org_id filter if provided if organization_id and organization_id not in org_admin_org_ids: raise HTTPException( status_code=403, detail={"error": "You can only view teams within your organizations."}, ) - # When the caller is an org admin querying their own teams (or no - # specific user), null out user_id so that - # _build_team_list_where_conditions scopes only by organization_id - # — org admins should see all teams in their orgs, not just teams - # they are a direct member of. Keep user_id when the org admin - # explicitly queries a *different* user's teams. - if user_id is None or user_id == user_api_key_dict.user_id: - user_id = None + is_own_query: Final = user_id is None or user_id == caller_user_id + own_team_ids: Final = ( + await _get_user_team_ids( + user_id=caller_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + if is_own_query + else () + ) verbose_proxy_logger.debug( "list_team_v2: org admin access for user=%s, org_ids=%s, user_id_filter=%s", - user_api_key_dict.user_id, + caller_user_id, org_admin_org_ids, - user_id, + None if is_own_query else user_id, ) - else: - # Not an org admin — fall back to standard route check - if not allowed_route_check_inside_route(user_api_key_dict=user_api_key_dict, requested_user_id=user_id): - raise HTTPException( - status_code=401, - detail={ - "error": f"Only admin users can query all teams/other teams. Your user role={user_api_key_dict.user_role}" - }, - ) - # Regular user — auto-inject caller's user_id - if user_id is None: - user_id = user_api_key_dict.user_id + return None if is_own_query else user_id, org_admin_org_ids, own_team_ids - return user_id, org_admin_org_ids + # Not an org admin — fall back to standard route check + if not allowed_route_check_inside_route(user_api_key_dict=user_api_key_dict, requested_user_id=user_id): + raise HTTPException( + status_code=401, + detail={ + "error": f"Only admin users can query all teams/other teams. Your user role={user_api_key_dict.user_role}" + }, + ) + # Regular user — auto-inject caller's user_id + return user_id if user_id is not None else caller_user_id, None, () @router.get( @@ -5163,7 +5198,7 @@ async def list_team_v2( ) # --- Access control --- - user_id, org_admin_org_ids = await _enforce_list_team_v2_access( + user_id, org_admin_org_ids, own_team_ids = await _enforce_list_team_v2_access( user_api_key_dict=user_api_key_dict, user_id=user_id, organization_id=organization_id, @@ -5195,6 +5230,7 @@ async def list_team_v2( search=search, search_team_id_match=search_team_id_match, org_admin_org_ids=org_admin_org_ids, + own_team_ids=own_team_ids, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) @@ -5290,18 +5326,18 @@ async def _authorize_and_filter_teams( Authorize the /team/list request and return filtered teams. - Proxy admins: all teams (or filtered by user_id if provided). - - Org admins: teams from their orgs (scoped to user_id if provided). + - Org admins: teams from their orgs (scoped to user_id if provided), plus + the teams they are a member of when querying themselves. - Own query (user_id matches caller): teams the user is a member of. - Others: 401. """ is_proxy_admin: Final = _user_has_admin_view(user_api_key_dict) + is_own_query: Final = ( + user_id is not None and user_api_key_dict.user_id is not None and user_api_key_dict.user_id == user_id + ) allowed_org_ids: list[str] | None = None if not is_proxy_admin: - is_own_query: Final = ( - user_id is not None and user_api_key_dict.user_id is not None and user_api_key_dict.user_id == user_id - ) - # Check if user is an org admin (even for own queries, so they see org teams) if user_api_key_dict.user_id is not None: caller_user: Final = await get_user_object( @@ -5328,33 +5364,30 @@ async def _authorize_and_filter_teams( }, ) - if allowed_org_ids is not None: - # Org admin: query DB for teams in their orgs + if allowed_org_ids is not None and user_id and not is_own_query: org_teams: Final = await _raw_team_db(TeamRepository(prisma_client)).find_many( where={"organization_id": {"in": allowed_org_ids}}, include={"litellm_model_table": True}, ) - if not user_id: - return list(org_teams) - # Filter org teams to only those where the target user is a member return [ team for team in org_teams if team.members_with_roles and any(m.get("user_id") == user_id for m in team.members_with_roles) ] - elif user_id: - # Regular user: fetch all and filter by membership (Prisma can't filter JSON arrays) - response: Final = await _raw_team_db(TeamRepository(prisma_client)).find_many( - include={"litellm_model_table": True} - ) - return [ - team - for team in response - if team.members_with_roles and any(m.get("user_id") == user_id for m in team.members_with_roles) - ] - else: + + response: Final = await _raw_team_db(TeamRepository(prisma_client)).find_many(include={"litellm_model_table": True}) + if allowed_org_ids is None and not user_id: # Proxy admin: all teams - return list(await _raw_team_db(TeamRepository(prisma_client)).find_many(include={"litellm_model_table": True})) + return list(response) + + # Prisma can't filter JSON arrays, so membership is filtered in Python + viewer_id: Final = user_id or user_api_key_dict.user_id + return [ + team + for team in response + if (allowed_org_ids is not None and team.organization_id in allowed_org_ids) + or (team.members_with_roles and any(m.get("user_id") == viewer_id for m in team.members_with_roles)) + ] @router.get("/team/list", tags=["team management"], dependencies=[Depends(user_api_key_auth)]) diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 0e1831614ac..0f891c10d87 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -4055,10 +4055,167 @@ async def test_list_team_v2_org_admin_own_user_id_sees_all_org_teams(): assert result["total"] == 2 assert len(result["teams"]) == 2 - # Verify the where clause scopes by org only — no team_id filter + # Verify the where clause scopes by org OR own membership — no + # top-level team_id filter that would hide org teams they aren't in where = mock_db.litellm_teamtable.find_many.call_args.kwargs["where"] - assert where["organization_id"] == {"in": ["org_A"]} + assert where["AND"] == [ + {"OR": [{"organization_id": {"in": ["org_A"]}}, {"team_id": {"in": ["team_1"]}}]} + ] assert "team_id" not in where + assert "organization_id" not in where + + +@pytest.mark.asyncio +async def test_list_team_v2_org_admin_own_query_keeps_memberships_in_other_orgs(monkeypatch): + """ + An org admin of org_A who is only a member of a team in org_B must still + see that team when listing their own teams: the where clause must union + org scope with membership instead of intersecting them. + + Regression test for LIT-3723. + """ + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.management_endpoints.team_endpoints import list_team_v2 + + org_admin = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="org_admin_user") + cache = UserApiKeyCache() + await cache.async_set_cache( + key="org_admin_user", + value=LiteLLM_UserTable( + user_id="org_admin_user", + teams=["team_in_org_A", "team_in_org_B"], + organization_memberships=[ + LiteLLM_OrganizationMembershipTable( + user_id="org_admin_user", + organization_id="org_A", + user_role="org_admin", + spend=0.0, + created_at=datetime.now(), + updated_at=datetime.now(), + ), + LiteLLM_OrganizationMembershipTable( + user_id="org_admin_user", + organization_id="org_B", + user_role="internal_user", + spend=0.0, + created_at=datetime.now(), + updated_at=datetime.now(), + ), + ], + ), + model_type=LiteLLM_UserTable, + ) + prisma_client = MagicMock() + prisma_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[]) + prisma_client.db.litellm_teamtable.count = AsyncMock(return_value=0) + prisma_client.db.litellm_verificationtoken.group_by = AsyncMock(return_value=[]) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache) + proxy_logging_obj = MagicMock() + proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj) + + for own_user_id in ("org_admin_user", None): + await list_team_v2( + http_request=MagicMock(), + user_id=own_user_id, + organization_id=None, + team_id=None, + team_alias=None, + search="team", + user_api_key_dict=org_admin, + page=1, + page_size=10, + sort_by=None, + sort_order="asc", + status=None, + ) + + where = prisma_client.db.litellm_teamtable.find_many.call_args.kwargs["where"] + assert where["AND"] == [ + { + "OR": [ + {"organization_id": {"in": ["org_A"]}}, + {"team_id": {"in": ["team_in_org_A", "team_in_org_B"]}}, + ] + } + ] + assert where["OR"] == [ + {"team_id": "team"}, + {"team_alias": {"contains": "team", "mode": "insensitive"}}, + ] + assert "organization_id" not in where + assert "team_id" not in where + + +@pytest.mark.asyncio +async def test_list_team_v1_org_admin_own_query_keeps_memberships_in_other_orgs(): + """ + /team/list: an org admin of org_A listing their own teams sees every team + in org_A plus the org_B team they are a member of, but a query for another + user stays scoped to org_A. + + Regression test for LIT-3723. + """ + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.management_endpoints.team_endpoints import _authorize_and_filter_teams + + org_admin = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="org_admin_user") + cache = UserApiKeyCache() + await cache.async_set_cache( + key="org_admin_user", + value=LiteLLM_UserTable( + user_id="org_admin_user", + teams=["team_in_org_A", "team_in_org_B"], + organization_memberships=[ + LiteLLM_OrganizationMembershipTable( + user_id="org_admin_user", + organization_id="org_A", + user_role="org_admin", + spend=0.0, + created_at=datetime.now(), + updated_at=datetime.now(), + ), + ], + ), + model_type=LiteLLM_UserTable, + ) + + def team(team_id, organization_id, *member_ids): + return SimpleNamespace( + team_id=team_id, + organization_id=organization_id, + members_with_roles=[{"user_id": m, "role": "user"} for m in member_ids], + ) + + all_teams = [ + team("team_in_org_A", "org_A", "org_admin_user"), + team("other_team_in_org_A", "org_A", "other_user"), + team("team_in_org_B", "org_B", "org_admin_user", "other_user"), + team("unrelated_team_in_org_B", "org_B", "other_user"), + ] + + async def find_many(where=None, **kwargs): + if where is None: + return all_teams + return [t for t in all_teams if t.organization_id in where["organization_id"]["in"]] + + prisma_client = MagicMock() + prisma_client.db.litellm_teamtable.find_many = AsyncMock(side_effect=find_many) + + async def list_teams(user_id): + teams = await _authorize_and_filter_teams( + user_api_key_dict=org_admin, + user_id=user_id, + prisma_client=prisma_client, + user_api_key_cache=cache, + proxy_logging_obj=MagicMock(), + ) + return [t.team_id for t in teams] + + assert await list_teams("org_admin_user") == ["team_in_org_A", "other_team_in_org_A", "team_in_org_B"] + assert await list_teams(None) == ["team_in_org_A", "other_team_in_org_A", "team_in_org_B"] + assert await list_teams("other_user") == ["other_team_in_org_A"] @pytest.mark.asyncio From 8d9edfc03c5051265f5d36fc1ce62814a990072c Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 15:40:41 +0000 Subject: [PATCH 036/112] fix(proxy): keep /team/list self query for org admins membership-only across orgs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/management_endpoints/team_endpoints.py | 15 +++++++-------- .../management_endpoints/test_team_endpoints.py | 8 ++++---- 2 files changed, 11 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index b4608ac27a0..4234206cd33 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -5326,9 +5326,8 @@ async def _authorize_and_filter_teams( Authorize the /team/list request and return filtered teams. - Proxy admins: all teams (or filtered by user_id if provided). - - Org admins: teams from their orgs (scoped to user_id if provided), plus - the teams they are a member of when querying themselves. - - Own query (user_id matches caller): teams the user is a member of. + - Org admins: teams from their orgs (scoped to user_id if provided). + - Own query (user_id matches caller): teams the user is a member of, across all orgs. - Others: 401. """ is_proxy_admin: Final = _user_has_admin_view(user_api_key_dict) @@ -5364,11 +5363,13 @@ async def _authorize_and_filter_teams( }, ) - if allowed_org_ids is not None and user_id and not is_own_query: + if allowed_org_ids is not None and not is_own_query: org_teams: Final = await _raw_team_db(TeamRepository(prisma_client)).find_many( where={"organization_id": {"in": allowed_org_ids}}, include={"litellm_model_table": True}, ) + if not user_id: + return list(org_teams) return [ team for team in org_teams @@ -5376,17 +5377,15 @@ async def _authorize_and_filter_teams( ] response: Final = await _raw_team_db(TeamRepository(prisma_client)).find_many(include={"litellm_model_table": True}) - if allowed_org_ids is None and not user_id: + if not user_id: # Proxy admin: all teams return list(response) # Prisma can't filter JSON arrays, so membership is filtered in Python - viewer_id: Final = user_id or user_api_key_dict.user_id return [ team for team in response - if (allowed_org_ids is not None and team.organization_id in allowed_org_ids) - or (team.members_with_roles and any(m.get("user_id") == viewer_id for m in team.members_with_roles)) + if team.members_with_roles and any(m.get("user_id") == user_id for m in team.members_with_roles) ] diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 0f891c10d87..7aaf558e2ba 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -4152,8 +4152,8 @@ async def test_list_team_v2_org_admin_own_query_keeps_memberships_in_other_orgs( async def test_list_team_v1_org_admin_own_query_keeps_memberships_in_other_orgs(): """ /team/list: an org admin of org_A listing their own teams sees every team - in org_A plus the org_B team they are a member of, but a query for another - user stays scoped to org_A. + they belong to, including the org_B one. The bare admin listing stays the + org_A view and a query for another user stays scoped to org_A. Regression test for LIT-3723. """ @@ -4213,8 +4213,8 @@ async def test_list_team_v1_org_admin_own_query_keeps_memberships_in_other_orgs( ) return [t.team_id for t in teams] - assert await list_teams("org_admin_user") == ["team_in_org_A", "other_team_in_org_A", "team_in_org_B"] - assert await list_teams(None) == ["team_in_org_A", "other_team_in_org_A", "team_in_org_B"] + assert await list_teams("org_admin_user") == ["team_in_org_A", "team_in_org_B"] + assert await list_teams(None) == ["team_in_org_A", "other_team_in_org_A"] assert await list_teams("other_user") == ["other_team_in_org_A"] From d64390978932babd0b08c8d2bfb3a191ad8377d0 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 16:03:23 +0000 Subject: [PATCH 037/112] fix(proxy): read org admin's own team ids from db and assert v2 list results Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints/team_endpoints.py | 27 +--- .../test_team_endpoints.py | 141 +++++++++++------- 2 files changed, 94 insertions(+), 74 deletions(-) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 4234206cd33..1a562f26776 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -4858,23 +4858,9 @@ async def _get_org_admin_org_ids( return org_ids if org_ids else None -async def _get_user_team_ids( - user_id: str, - prisma_client: PrismaClient, - user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: ProxyLogging, -) -> tuple[str, ...]: - try: - user: Final = await get_user_object( - user_id=user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - user_id_upsert=False, - proxy_logging_obj=proxy_logging_obj, - ) - except ValueError: - return () - return tuple(user.teams or ()) if user is not None else () +async def _get_user_team_ids_from_db(user_id: str, prisma_client: PrismaClient) -> tuple[str, ...]: + user_row: Final = await _user_db(prisma_client).find_unique(where={"user_id": user_id}) + return tuple(user_row.teams or ()) if user_row is not None else () async def _build_team_list_where_conditions( @@ -5094,12 +5080,7 @@ async def _enforce_list_team_v2_access( ) is_own_query: Final = user_id is None or user_id == caller_user_id own_team_ids: Final = ( - await _get_user_team_ids( - user_id=caller_user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) + await _get_user_team_ids_from_db(user_id=caller_user_id, prisma_client=prisma_client) if is_own_query else () ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 7aaf558e2ba..9a3d85d84ff 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -3937,6 +3937,7 @@ async def test_list_team_v2_org_admin_sees_org_teams(): mock_db.litellm_teamtable.find_many = AsyncMock(return_value=[mock_team]) mock_db.litellm_teamtable.count = AsyncMock(return_value=1) mock_db.litellm_verificationtoken.group_by = AsyncMock(return_value=[]) + mock_db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user) result = await list_team_v2( http_request=mock_request, @@ -4036,6 +4037,7 @@ async def test_list_team_v2_org_admin_own_user_id_sees_all_org_teams(): ) mock_db.litellm_teamtable.count = AsyncMock(return_value=2) mock_db.litellm_verificationtoken.group_by = AsyncMock(return_value=[]) + mock_db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user) # UI sends the caller's own user_id for non-Admin roles result = await list_team_v2( @@ -4065,12 +4067,45 @@ async def test_list_team_v2_org_admin_own_user_id_sees_all_org_teams(): assert "organization_id" not in where +def _team_where_matches(team, where) -> bool: + for key, cond in where.items(): + if key == "AND": + if not all(_team_where_matches(team, c) for c in cond): + return False + elif key == "OR": + if not any(_team_where_matches(team, c) for c in cond): + return False + else: + value = getattr(team, key) + if not isinstance(cond, dict): + if value != cond: + return False + elif "in" in cond and value not in cond["in"]: + return False + elif "contains" in cond and cond["contains"].lower() not in (value or "").lower(): + return False + return True + + +def _org_membership(user_id: str, organization_id: str, user_role: str) -> LiteLLM_OrganizationMembershipTable: + return LiteLLM_OrganizationMembershipTable( + user_id=user_id, + organization_id=organization_id, + user_role=user_role, + spend=0.0, + created_at=datetime.now(), + updated_at=datetime.now(), + ) + + @pytest.mark.asyncio async def test_list_team_v2_org_admin_own_query_keeps_memberships_in_other_orgs(monkeypatch): """ - An org admin of org_A who is only a member of a team in org_B must still - see that team when listing their own teams: the where clause must union - org scope with membership instead of intersecting them. + /v2/team/list: an org admin of org_A who is a member of a team in org_B + gets that team back on a self query (with and without user_id, with and + without search), alongside every org_A team. The membership half of the + union comes from the DB, so a stale cached user object cannot hide it. + A query for another user stays scoped to org_A. Regression test for LIT-3723. """ @@ -4083,46 +4118,66 @@ async def test_list_team_v2_org_admin_own_query_keeps_memberships_in_other_orgs( key="org_admin_user", value=LiteLLM_UserTable( user_id="org_admin_user", - teams=["team_in_org_A", "team_in_org_B"], + teams=["team_in_org_A"], organization_memberships=[ - LiteLLM_OrganizationMembershipTable( - user_id="org_admin_user", - organization_id="org_A", - user_role="org_admin", - spend=0.0, - created_at=datetime.now(), - updated_at=datetime.now(), - ), - LiteLLM_OrganizationMembershipTable( - user_id="org_admin_user", - organization_id="org_B", - user_role="internal_user", - spend=0.0, - created_at=datetime.now(), - updated_at=datetime.now(), - ), + _org_membership("org_admin_user", "org_A", "org_admin"), + _org_membership("org_admin_user", "org_B", "internal_user"), ], ), model_type=LiteLLM_UserTable, ) + await cache.async_set_cache( + key="other_user", + value=LiteLLM_UserTable( + user_id="other_user", + teams=["other_team_in_org_A", "team_in_org_B", "unrelated_team_in_org_B"], + organization_memberships=[_org_membership("other_user", "org_B", "internal_user")], + ), + model_type=LiteLLM_UserTable, + ) + + def team(team_id, organization_id, *member_ids): + return LiteLLM_TeamTable( + team_id=team_id, + team_alias=team_id, + organization_id=organization_id, + members_with_roles=[Member(user_id=m, role="user") for m in member_ids], + ) + + all_teams = [ + team("team_in_org_A", "org_A", "org_admin_user"), + team("other_team_in_org_A", "org_A", "other_user"), + team("team_in_org_B", "org_B", "org_admin_user", "other_user"), + team("unrelated_team_in_org_B", "org_B", "other_user"), + ] + + async def find_many(where=None, **kwargs): + return [t for t in all_teams if where is None or _team_where_matches(t, where)] + + async def count(where=None, **kwargs): + return len(await find_many(where)) + prisma_client = MagicMock() - prisma_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[]) - prisma_client.db.litellm_teamtable.count = AsyncMock(return_value=0) + prisma_client.db.litellm_teamtable.find_many = AsyncMock(side_effect=find_many) + prisma_client.db.litellm_teamtable.count = AsyncMock(side_effect=count) prisma_client.db.litellm_verificationtoken.group_by = AsyncMock(return_value=[]) + prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=SimpleNamespace(teams=["team_in_org_A", "team_in_org_B"]) + ) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client) monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache) proxy_logging_obj = MagicMock() proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock() monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj) - for own_user_id in ("org_admin_user", None): - await list_team_v2( + async def list_teams(user_id, search=None): + result = await list_team_v2( http_request=MagicMock(), - user_id=own_user_id, + user_id=user_id, organization_id=None, team_id=None, team_alias=None, - search="team", + search=search, user_api_key_dict=org_admin, page=1, page_size=10, @@ -4130,22 +4185,15 @@ async def test_list_team_v2_org_admin_own_query_keeps_memberships_in_other_orgs( sort_order="asc", status=None, ) + assert result["total"] == len(result["teams"]) + return [t.team_id for t in result["teams"]] - where = prisma_client.db.litellm_teamtable.find_many.call_args.kwargs["where"] - assert where["AND"] == [ - { - "OR": [ - {"organization_id": {"in": ["org_A"]}}, - {"team_id": {"in": ["team_in_org_A", "team_in_org_B"]}}, - ] - } - ] - assert where["OR"] == [ - {"team_id": "team"}, - {"team_alias": {"contains": "team", "mode": "insensitive"}}, - ] - assert "organization_id" not in where - assert "team_id" not in where + own_view = ["team_in_org_A", "other_team_in_org_A", "team_in_org_B"] + assert await list_teams("org_admin_user") == own_view + assert await list_teams(None) == own_view + assert await list_teams("org_admin_user", search="team_in_org_B") == ["team_in_org_B"] + assert await list_teams("other_user") == ["other_team_in_org_A"] + prisma_client.db.litellm_usertable.find_unique.assert_awaited_with(where={"user_id": "org_admin_user"}) @pytest.mark.asyncio @@ -4167,16 +4215,7 @@ async def test_list_team_v1_org_admin_own_query_keeps_memberships_in_other_orgs( value=LiteLLM_UserTable( user_id="org_admin_user", teams=["team_in_org_A", "team_in_org_B"], - organization_memberships=[ - LiteLLM_OrganizationMembershipTable( - user_id="org_admin_user", - organization_id="org_A", - user_role="org_admin", - spend=0.0, - created_at=datetime.now(), - updated_at=datetime.now(), - ), - ], + organization_memberships=[_org_membership("org_admin_user", "org_A", "org_admin")], ), model_type=LiteLLM_UserTable, ) From dc30be006bbc8cf9df8801780814fccdab6f491e Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 16:27:31 +0000 Subject: [PATCH 038/112] refactor(proxy): read org admin's own team ids via get_user_object(check_db_only=True) Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints/team_endpoints.py | 28 ++++++++++++++++--- .../test_team_endpoints.py | 13 +++++++-- 2 files changed, 35 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 1a562f26776..3a77f163251 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -4858,9 +4858,24 @@ async def _get_org_admin_org_ids( return org_ids if org_ids else None -async def _get_user_team_ids_from_db(user_id: str, prisma_client: PrismaClient) -> tuple[str, ...]: - user_row: Final = await _user_db(prisma_client).find_unique(where={"user_id": user_id}) - return tuple(user_row.teams or ()) if user_row is not None else () +async def _get_user_team_ids_from_db( + user_id: str, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> tuple[str, ...]: + try: + user: Final = await get_user_object( + user_id=user_id, + 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, + ) + except ValueError: + return () + return tuple(user.teams or ()) if user is not None else () async def _build_team_list_where_conditions( @@ -5080,7 +5095,12 @@ async def _enforce_list_team_v2_access( ) is_own_query: Final = user_id is None or user_id == caller_user_id own_team_ids: Final = ( - await _get_user_team_ids_from_db(user_id=caller_user_id, prisma_client=prisma_client) + await _get_user_team_ids_from_db( + user_id=caller_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) if is_own_query else () ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 9a3d85d84ff..82137f7861e 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -4162,7 +4162,14 @@ async def test_list_team_v2_org_admin_own_query_keeps_memberships_in_other_orgs( prisma_client.db.litellm_teamtable.count = AsyncMock(side_effect=count) prisma_client.db.litellm_verificationtoken.group_by = AsyncMock(return_value=[]) prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=SimpleNamespace(teams=["team_in_org_A", "team_in_org_B"]) + return_value=LiteLLM_UserTable( + user_id="org_admin_user", + teams=["team_in_org_A", "team_in_org_B"], + organization_memberships=[ + _org_membership("org_admin_user", "org_A", "org_admin"), + _org_membership("org_admin_user", "org_B", "internal_user"), + ], + ) ) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client) monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache) @@ -4193,7 +4200,9 @@ async def test_list_team_v2_org_admin_own_query_keeps_memberships_in_other_orgs( assert await list_teams(None) == own_view assert await list_teams("org_admin_user", search="team_in_org_B") == ["team_in_org_B"] assert await list_teams("other_user") == ["other_team_in_org_A"] - prisma_client.db.litellm_usertable.find_unique.assert_awaited_with(where={"user_id": "org_admin_user"}) + prisma_client.db.litellm_usertable.find_unique.assert_awaited_with( + where={"user_id": "org_admin_user"}, include={"organization_memberships": True} + ) @pytest.mark.asyncio From 2268bbaf5ef0ac5a70bc75f77e2dc0203f4b0eef Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 14 Sep 2026 16:39:54 +0000 Subject: [PATCH 039/112] feat(proxy): honor LITELLM_DISABLE_ACCESS_LOG_PATHS to drop noisy uvicorn access log lines Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/_logging.py | 32 +++++++++++++++ tests/test_litellm/test_logging.py | 63 ++++++++++++++++++++++++++++++ 2 files changed, 95 insertions(+) diff --git a/litellm/_logging.py b/litellm/_logging.py index c73b5175a31..33bd7b7bbf3 100644 --- a/litellm/_logging.py +++ b/litellm/_logging.py @@ -225,6 +225,37 @@ class AccessLogRedactionFilter(logging.Filter): _access_log_filter: Final = AccessLogRedactionFilter() +def _parse_disabled_access_log_paths(raw: str) -> frozenset[str]: + return frozenset(path for path in (segment.strip() for segment in raw.split(",")) if path) + + +_DISABLED_ACCESS_LOG_PATHS: Final = _parse_disabled_access_log_paths(os.getenv("LITELLM_DISABLE_ACCESS_LOG_PATHS", "")) + + +class AccessLogPathFilter(logging.Filter): + """Drops uvicorn.access records for request paths listed in LITELLM_DISABLE_ACCESS_LOG_PATHS. + + uvicorn passes record.args as (client_addr, method, full_path, http_version, status_code). + """ + + def __init__(self, disabled_paths: frozenset[str]) -> None: + super().__init__() + self._disabled_paths: Final = disabled_paths + + def filter(self, record: logging.LogRecord) -> bool: + if not self._disabled_paths: + return True + if not isinstance(record.args, tuple) or len(record.args) < 3: + return True + full_path: Final = record.args[2] + if not isinstance(full_path, str): + return True + return full_path.partition("?")[0] not in self._disabled_paths + + +_access_log_path_filter: Final = AccessLogPathFilter(_DISABLED_ACCESS_LOG_PATHS) + + def _get_max_string_length_stdout_log() -> int: """Read the limit per record so a value loaded later via proxy config environment_variables is honored.""" @@ -664,6 +695,7 @@ def _redact_third_party_loggers() -> None: logging.getLogger(name).addFilter(_secret_filter) for name in _REDACTED_ACCESS_LOGGERS: logging.getLogger(name).addFilter(_access_log_filter) + logging.getLogger(name).addFilter(_access_log_path_filter) # Call the suppression function diff --git a/tests/test_litellm/test_logging.py b/tests/test_litellm/test_logging.py index bf3757e6886..9972be01f00 100644 --- a/tests/test_litellm/test_logging.py +++ b/tests/test_litellm/test_logging.py @@ -2,6 +2,7 @@ import ast import asyncio import json import logging +import os import re import sys import time @@ -17,6 +18,7 @@ from litellm._logging import ( _MAX_SCRUBBED_ACCESS_ARG, _PLAIN_LOG_FORMAT, ALL_LOGGERS, + AccessLogPathFilter, AccessLogRedactionFilter, CorrelationContextFilter, CorrelationPlainFormatter, @@ -26,6 +28,7 @@ from litellm._logging import ( StdoutLogTruncationFilter, _get_uvicorn_json_log_config, _initialize_loggers_with_handler, + _parse_disabled_access_log_paths, _parse_json_logs_env, _plain_log_format, _stdout_truncation_marker, @@ -1178,3 +1181,63 @@ def test_access_redaction_survives_the_uvicorn_json_log_config(): lg.handlers[:] = handlers lg.setLevel(level) lg.propagate = True + + +@pytest.mark.parametrize( + "raw, expected", + [ + ("", frozenset()), + ("/health,/metrics", frozenset({"/health", "/metrics"})), + (" /health , /metrics/ ,", frozenset({"/health", "/metrics/"})), + ], +) +def test_parse_disabled_access_log_paths(raw, expected): + assert _parse_disabled_access_log_paths(raw) == expected + + +def test_access_log_path_filter_drops_a_listed_path(): + access_log_filter = AccessLogPathFilter(frozenset({"/health/liveliness", "/metrics"})) + + assert access_log_filter.filter(_access_record("/health/liveliness")) is False + assert access_log_filter.filter(_access_record("/metrics?format=prometheus")) is False + + +def test_access_log_path_filter_keeps_an_unlisted_path(): + access_log_filter = AccessLogPathFilter(frozenset({"/health/liveliness", "/metrics"})) + + assert access_log_filter.filter(_access_record("/v1/chat/completions")) is True + assert access_log_filter.filter(_access_record("/health")) is True + assert access_log_filter.filter(_access_record("/metrics/")) is True + + +def test_access_log_path_filter_is_a_no_op_when_unset(): + assert AccessLogPathFilter(frozenset()).filter(_access_record("/health/liveliness")) is True + + +def test_access_log_path_filter_keeps_a_record_without_positional_args(): + record = logging.LogRecord( + name="uvicorn.access", + level=logging.INFO, + pathname="", + lineno=0, + msg="access line", + args=None, + exc_info=None, + ) + + assert AccessLogPathFilter(frozenset({"/health/liveliness"})).filter(record) is True + + +def test_uvicorn_access_logger_drops_a_listed_path_end_to_end(monkeypatch): + monkeypatch.setenv("LITELLM_DISABLE_ACCESS_LOG_PATHS", "/health/liveliness") + access_log_filter = AccessLogPathFilter( + _parse_disabled_access_log_paths(os.environ["LITELLM_DISABLE_ACCESS_LOG_PATHS"]) + ) + logger = logging.getLogger("uvicorn.access") + assert any(isinstance(f, AccessLogPathFilter) for f in logger.filters) + logger.addFilter(access_log_filter) + try: + assert _emit_access_line("/health/liveliness") == "" + assert _emit_access_line("/v1/chat/completions") != "" + finally: + logger.removeFilter(access_log_filter) From 4b39344d137444cb153a71660ab901d05efe31b5 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 16:49:54 +0000 Subject: [PATCH 040/112] fix(proxy): only treat a missing user as no memberships on team list, surface db errors Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/management_endpoints/team_endpoints.py | 3 ++- .../proxy/management_endpoints/test_team_endpoints.py | 4 ++++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 3a77f163251..fdd44c16799 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -76,6 +76,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.auth_checks import ( OrganizationNotFoundError, + UserNotFoundError, _cache_team_object, allowed_route_check_inside_route, can_org_access_model, @@ -4873,7 +4874,7 @@ async def _get_user_team_ids_from_db( proxy_logging_obj=proxy_logging_obj, check_db_only=True, ) - except ValueError: + except UserNotFoundError: return () return tuple(user.teams or ()) if user is not None else () diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 82137f7861e..ccb66fd9534 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -4204,6 +4204,10 @@ async def test_list_team_v2_org_admin_own_query_keeps_memberships_in_other_orgs( where={"user_id": "org_admin_user"}, include={"organization_memberships": True} ) + prisma_client.db.litellm_usertable.find_unique.side_effect = RuntimeError("db down") + with pytest.raises(ValueError, match="db down"): + await list_teams("org_admin_user") + @pytest.mark.asyncio async def test_list_team_v1_org_admin_own_query_keeps_memberships_in_other_orgs(): From 66ce1eea97e32da969f5268dc4d0b75fc7929d6c Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 14 Sep 2026 16:54:38 +0000 Subject: [PATCH 041/112] test(proxy): cover non-string access log paths Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/test_logging.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/tests/test_litellm/test_logging.py b/tests/test_litellm/test_logging.py index 9972be01f00..7106d25b7c2 100644 --- a/tests/test_litellm/test_logging.py +++ b/tests/test_litellm/test_logging.py @@ -1228,6 +1228,20 @@ def test_access_log_path_filter_keeps_a_record_without_positional_args(): assert AccessLogPathFilter(frozenset({"/health/liveliness"})).filter(record) is True +def test_access_log_path_filter_keeps_a_record_with_a_non_string_path(): + record = logging.LogRecord( + name="uvicorn.access", + level=logging.INFO, + pathname="", + lineno=0, + msg="%s - %s %s", + args=("127.0.0.1:1", "GET", 42), + exc_info=None, + ) + + assert AccessLogPathFilter(frozenset({"/health/liveliness"})).filter(record) is True + + def test_uvicorn_access_logger_drops_a_listed_path_end_to_end(monkeypatch): monkeypatch.setenv("LITELLM_DISABLE_ACCESS_LOG_PATHS", "/health/liveliness") access_log_filter = AccessLogPathFilter( From 83d16a4690ad1d22f26b9df36b82f000bbd1270b Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 17:27:43 +0000 Subject: [PATCH 042/112] fix(proxy): read LITELLM_DISABLE_ACCESS_LOG_PATHS per record and match before redaction Values loaded after import via proxy config environment_variables or dotenv were ignored, and a long query string was truncated by the redaction filter before the path filter could match it. Tests now go through the production registration on the uvicorn.access logger instead of a hand-built filter. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/_logging.py | 21 +++--- tests/test_litellm/test_logging.py | 115 ++++++++++++++--------------- 2 files changed, 64 insertions(+), 72 deletions(-) diff --git a/litellm/_logging.py b/litellm/_logging.py index 33bd7b7bbf3..03a9bcf21cf 100644 --- a/litellm/_logging.py +++ b/litellm/_logging.py @@ -1,5 +1,6 @@ import ast import contextvars +import functools import logging import os import sys @@ -225,11 +226,15 @@ class AccessLogRedactionFilter(logging.Filter): _access_log_filter: Final = AccessLogRedactionFilter() +@functools.lru_cache(maxsize=1) def _parse_disabled_access_log_paths(raw: str) -> frozenset[str]: - return frozenset(path for path in (segment.strip() for segment in raw.split(",")) if path) + return frozenset(stripped for path in raw.split(",") if (stripped := path.strip())) -_DISABLED_ACCESS_LOG_PATHS: Final = _parse_disabled_access_log_paths(os.getenv("LITELLM_DISABLE_ACCESS_LOG_PATHS", "")) +def _disabled_access_log_paths() -> frozenset[str]: + """Read the variable per record so a value loaded later via proxy config + environment_variables or dotenv is honored.""" + return _parse_disabled_access_log_paths(os.getenv("LITELLM_DISABLE_ACCESS_LOG_PATHS", "")) class AccessLogPathFilter(logging.Filter): @@ -238,22 +243,16 @@ class AccessLogPathFilter(logging.Filter): uvicorn passes record.args as (client_addr, method, full_path, http_version, status_code). """ - def __init__(self, disabled_paths: frozenset[str]) -> None: - super().__init__() - self._disabled_paths: Final = disabled_paths - def filter(self, record: logging.LogRecord) -> bool: - if not self._disabled_paths: - return True if not isinstance(record.args, tuple) or len(record.args) < 3: return True full_path: Final = record.args[2] if not isinstance(full_path, str): return True - return full_path.partition("?")[0] not in self._disabled_paths + return full_path.partition("?")[0] not in _disabled_access_log_paths() -_access_log_path_filter: Final = AccessLogPathFilter(_DISABLED_ACCESS_LOG_PATHS) +_access_log_path_filter: Final = AccessLogPathFilter() def _get_max_string_length_stdout_log() -> int: @@ -694,8 +693,8 @@ def _redact_third_party_loggers() -> None: for name in _REDACTED_THIRD_PARTY_LOGGERS: logging.getLogger(name).addFilter(_secret_filter) for name in _REDACTED_ACCESS_LOGGERS: - logging.getLogger(name).addFilter(_access_log_filter) logging.getLogger(name).addFilter(_access_log_path_filter) + logging.getLogger(name).addFilter(_access_log_filter) # Call the suppression function diff --git a/tests/test_litellm/test_logging.py b/tests/test_litellm/test_logging.py index 7106d25b7c2..57c39280a9f 100644 --- a/tests/test_litellm/test_logging.py +++ b/tests/test_litellm/test_logging.py @@ -2,7 +2,6 @@ import ast import asyncio import json import logging -import os import re import sys import time @@ -28,7 +27,6 @@ from litellm._logging import ( StdoutLogTruncationFilter, _get_uvicorn_json_log_config, _initialize_loggers_with_handler, - _parse_disabled_access_log_paths, _parse_json_logs_env, _plain_log_format, _stdout_truncation_marker, @@ -1183,75 +1181,70 @@ def test_access_redaction_survives_the_uvicorn_json_log_config(): lg.propagate = True +_DISABLED_ACCESS_LOG_PATHS_RAW = " /health/liveliness , ,/metrics/" + + @pytest.mark.parametrize( - "raw, expected", + "full_path", [ - ("", frozenset()), - ("/health,/metrics", frozenset({"/health", "/metrics"})), - (" /health , /metrics/ ,", frozenset({"/health", "/metrics/"})), + "/health/liveliness", + "/health/liveliness?x=1", + "/health/liveliness?probe=" + "x" * _MAX_SCRUBBED_ACCESS_ARG, + "/metrics/", + "/metrics/?format=prometheus&job=a", ], ) -def test_parse_disabled_access_log_paths(raw, expected): - assert _parse_disabled_access_log_paths(raw) == expected +def test_uvicorn_access_logger_drops_a_configured_path(monkeypatch, full_path): + monkeypatch.setenv("LITELLM_DISABLE_ACCESS_LOG_PATHS", _DISABLED_ACCESS_LOG_PATHS_RAW) + assert _emit_access_line(full_path) == "" -def test_access_log_path_filter_drops_a_listed_path(): - access_log_filter = AccessLogPathFilter(frozenset({"/health/liveliness", "/metrics"})) - - assert access_log_filter.filter(_access_record("/health/liveliness")) is False - assert access_log_filter.filter(_access_record("/metrics?format=prometheus")) is False +@pytest.mark.parametrize( + "full_path", + ["/v1/chat/completions", "/health", "/health/liveliness/", "/metrics", "/v1/models?health=/health/liveliness"], +) +def test_uvicorn_access_logger_keeps_an_unconfigured_path(monkeypatch, full_path): + monkeypatch.setenv("LITELLM_DISABLE_ACCESS_LOG_PATHS", _DISABLED_ACCESS_LOG_PATHS_RAW) + assert f'"GET {full_path} HTTP/1.1" 200' in _emit_access_line(full_path) -def test_access_log_path_filter_keeps_an_unlisted_path(): - access_log_filter = AccessLogPathFilter(frozenset({"/health/liveliness", "/metrics"})) - - assert access_log_filter.filter(_access_record("/v1/chat/completions")) is True - assert access_log_filter.filter(_access_record("/health")) is True - assert access_log_filter.filter(_access_record("/metrics/")) is True +@pytest.mark.parametrize("raw", [None, "", " , ,"]) +def test_uvicorn_access_logger_keeps_every_line_when_no_path_is_configured(monkeypatch, raw): + if raw is None: + monkeypatch.delenv("LITELLM_DISABLE_ACCESS_LOG_PATHS", raising=False) + else: + monkeypatch.setenv("LITELLM_DISABLE_ACCESS_LOG_PATHS", raw) + assert '"GET /health/liveliness HTTP/1.1" 200' in _emit_access_line("/health/liveliness") -def test_access_log_path_filter_is_a_no_op_when_unset(): - assert AccessLogPathFilter(frozenset()).filter(_access_record("/health/liveliness")) is True +def test_access_log_path_filter_survives_the_uvicorn_json_log_config(monkeypatch): + import logging.config - -def test_access_log_path_filter_keeps_a_record_without_positional_args(): - record = logging.LogRecord( - name="uvicorn.access", - level=logging.INFO, - pathname="", - lineno=0, - msg="access line", - args=None, - exc_info=None, - ) - - assert AccessLogPathFilter(frozenset({"/health/liveliness"})).filter(record) is True - - -def test_access_log_path_filter_keeps_a_record_with_a_non_string_path(): - record = logging.LogRecord( - name="uvicorn.access", - level=logging.INFO, - pathname="", - lineno=0, - msg="%s - %s %s", - args=("127.0.0.1:1", "GET", 42), - exc_info=None, - ) - - assert AccessLogPathFilter(frozenset({"/health/liveliness"})).filter(record) is True - - -def test_uvicorn_access_logger_drops_a_listed_path_end_to_end(monkeypatch): - monkeypatch.setenv("LITELLM_DISABLE_ACCESS_LOG_PATHS", "/health/liveliness") - access_log_filter = AccessLogPathFilter( - _parse_disabled_access_log_paths(os.environ["LITELLM_DISABLE_ACCESS_LOG_PATHS"]) - ) - logger = logging.getLogger("uvicorn.access") - assert any(isinstance(f, AccessLogPathFilter) for f in logger.filters) - logger.addFilter(access_log_filter) + monkeypatch.setenv("LITELLM_DISABLE_ACCESS_LOG_PATHS", _DISABLED_ACCESS_LOG_PATHS_RAW) + names = ("uvicorn", "uvicorn.error", "uvicorn.access") + saved = tuple((logging.getLogger(n), logging.getLogger(n).handlers[:], logging.getLogger(n).level) for n in names) try: - assert _emit_access_line("/health/liveliness") == "" - assert _emit_access_line("/v1/chat/completions") != "" + logging.config.dictConfig(_get_uvicorn_json_log_config()) + + assert _emit_access_line("/health/liveliness?x=1") == "" + assert '"GET /v1/models HTTP/1.1" 200' in _emit_access_line("/v1/models") finally: - logger.removeFilter(access_log_filter) + for lg, handlers, level in saved: + lg.handlers[:] = handlers + lg.setLevel(level) + lg.propagate = True + + +@pytest.mark.parametrize("args", [None, ("127.0.0.1:1", "GET", 42)]) +def test_access_log_path_filter_keeps_a_record_without_a_string_path_arg(monkeypatch, args): + monkeypatch.setenv("LITELLM_DISABLE_ACCESS_LOG_PATHS", "/health/liveliness") + record = logging.LogRecord( + name="uvicorn.access", + level=logging.INFO, + pathname="", + lineno=0, + msg='127.0.0.1:1 - "GET /health/liveliness HTTP/1.1" 200', + args=args, + exc_info=None, + ) + assert AccessLogPathFilter().filter(record) is True From 34e702c5719e63ab51b557d758204377bc9faf6c Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 14 Sep 2026 17:46:18 +0000 Subject: [PATCH 043/112] test(otel): build the clipped-blob conversation without mutating the payload Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../integrations/otel/test_otel_v2_emitter.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py b/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py index 16fbb242ebd..6e2e467b856 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py @@ -518,9 +518,15 @@ def test_short_conversation_keeps_every_message_indexed(): def test_indexed_prompt_keeps_opener_and_latest_turns_under_a_value_length_limit(monkeypatch): """The system prompt and the live turn keep their own keys once the SDK clips ``input.value``.""" monkeypatch.setenv("OTEL_SPAN_ATTRIBUTE_VALUE_LENGTH_LIMIT", "256") - payload = _conversation_payload(60) - payload["messages"][0] = {"role": "system", "content": "be terse"} - payload["messages"][-1] = {"role": "user", "content": "LATEST-TURN"} + chat = _conversation_payload(60) + payload = { + **chat, + "messages": [ + {"role": "system", "content": "be terse"}, + *chat["messages"][1:-1], + {"role": "user", "content": "LATEST-TURN"}, + ], + } a = _conversation_span(["genai", "openinference"], payload).attributes assert len(a["input.value"]) == 256 From 5457f482905419ff966a237018cf2ce03fb21506 Mon Sep 17 00:00:00 2001 From: Shivi Jain Date: Mon, 14 Sep 2026 20:50:28 +0530 Subject: [PATCH 044/112] fix(auth): load team membership once per request and skip prisma on an L1 hit common_checks was querying get_team_membership twice, and DualCache awaited Redis SET on the auth path, so LRU eviction plus a hung Redis write showed up as two postgres spans --- litellm/proxy/auth/auth_checks.py | 313 ++++++++++++++---- .../proxy/auth/test_auth_checks.py | 127 +++++++ 2 files changed, 369 insertions(+), 71 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 576585ee9a3..abcade6b1c8 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -327,6 +327,9 @@ _safe_json_loads_obj: Final = _typed_json_loads(safe_json_loads) last_db_access_time: Final = LimitedSizeOrderedDict(max_size=100) db_cache_expiry: Final = DEFAULT_IN_MEMORY_TTL # refresh every 5s +_TEAM_MEMBERSHIP_CACHE_MISS: Final = object() +_team_membership_inflight: dict[str, asyncio.Task[LiteLLM_TeamMembership | None]] = {} + all_routes: Final = LiteLLMRoutes.openai_routes.value + LiteLLMRoutes.management_routes.value @@ -873,6 +876,21 @@ async def common_checks( """ from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + # One membership read for model-access, access-group attribution, and + # member-budget. Each used to call get_team_membership independently; + # DualCache Redis SET/GET on every miss made those look like two Postgres spans. + loaded_team_membership: LiteLLM_TeamMembership | None = None + team_membership_loaded = False + if team_object is not None and valid_token is not None and valid_token.user_id is not None: + loaded_team_membership = await get_team_membership( + user_id=valid_token.user_id, + team_id=team_object.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + team_membership_loaded = True + _model: Final[str | list[str] | None] = get_model_from_request( request_data=request_body, route=route, @@ -936,6 +954,8 @@ async def common_checks( prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, + team_membership=loaded_team_membership, + team_membership_loaded=team_membership_loaded, ) # Require trace id for agent keys when agent has require_trace_id_on_calls_by_agent @@ -987,6 +1007,8 @@ async def common_checks( prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, + team_membership=loaded_team_membership, + team_membership_loaded=team_membership_loaded, ) # Run before apply_key_tags_pre_auth injects key metadata.tags into request_body. @@ -1096,6 +1118,8 @@ async def common_checks( prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, + team_membership=loaded_team_membership, + team_membership_loaded=team_membership_loaded, ), _check_end_user_budget(end_user_obj=end_user_object, route=route) if end_user_object is not None and end_user_object.litellm_budget_table is not None @@ -2141,7 +2165,142 @@ async def get_tag_object( return tag_objects.get(tag_name) +def _debug_team_membership_log(hypothesis_id: str, message: str, data: dict[str, object]) -> None: + # #region agent log + try: + import json as _json + + with open("/Users/shijain/genai-apps/genai-proxy/.cursor/debug-86534f.log", "a", encoding="utf-8") as _f: + _f.write( + _json.dumps( + { + "sessionId": "86534f", + "timestamp": int(time.time() * 1000), + "location": "auth_checks.py:get_team_membership", + "message": message, + "hypothesisId": hypothesis_id, + "data": data, + } + ) + + "\n" + ) + except Exception: + pass + # #endregion + + +def _membership_from_cached_payload(cached: object) -> LiteLLM_TeamMembership | None | object: + """Decode a DualCache payload. ``_TEAM_MEMBERSHIP_CACHE_MISS`` means try the next tier.""" + if cached is None: + return _TEAM_MEMBERSHIP_CACHE_MISS + if cached == NO_TEAM_MEMBERSHIP_SENTINEL: + return None + cached_membership: Final = CacheCodec.deserialize(cached, model_type=LiteLLM_TeamMembership) + return cached_membership if cached_membership is not None else _TEAM_MEMBERSHIP_CACHE_MISS + + +async def _populate_team_membership_cache( + user_api_key_cache: UserApiKeyCache, + key: str, + value: object, + *, + model_type: type[LiteLLM_TeamMembership] | None = None, + ttl: float | None = None, +) -> None: + """Await in-memory write; replicate to Redis off the auth await path.""" + kwargs: dict[str, object] = {} + if model_type is not None: + kwargs["model_type"] = model_type + if ttl is not None: + kwargs["ttl"] = ttl + await user_api_key_cache.async_set_cache(key=key, value=value, local_only=True, **kwargs) + + async def _replicate_to_redis() -> None: + try: + await user_api_key_cache.async_set_cache(key=key, value=value, **kwargs) + except Exception: + return + + asyncio.create_task(_replicate_to_redis()) + + @log_db_metrics +async def _fetch_team_membership_from_db( + user_id: str, + team_id: str, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None = None, + proxy_logging_obj: ProxyLogging | None = None, +) -> LiteLLM_TeamMembership | None: + """Prisma read + L1 populate. Decorated so cache hits on ``get_team_membership`` are not postgres spans.""" + _ = parent_otel_span, proxy_logging_obj + response: Final = await _dictable_table(TeamMembershipRepository(prisma_client)).find_unique( + where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, + include={"litellm_budget_table": True}, + ) + _key: Final = team_membership_reservation_cache_key(user_id=user_id, team_id=team_id) + if response is None: + await _populate_team_membership_cache( + user_api_key_cache, + _key, + NO_TEAM_MEMBERSHIP_SENTINEL, + ttl=get_management_object_ttl(user_api_key_cache), + ) + return None + + membership: Final = LiteLLM_TeamMembership.model_validate(response.dict()) + await _populate_team_membership_cache( + user_api_key_cache, + _key, + membership, + model_type=LiteLLM_TeamMembership, + ) + return membership + + +async def _load_team_membership_on_cache_miss( + user_id: str, + team_id: str, + cache_key: str, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + proxy_logging_obj: ProxyLogging | None, +) -> LiteLLM_TeamMembership | None: + try: + redis_cached: Final[object] = await user_api_key_cache.async_get_cache(key=cache_key) + redis_membership: Final = _membership_from_cached_payload(redis_cached) + if redis_membership is not _TEAM_MEMBERSHIP_CACHE_MISS: + # #region agent log + _debug_team_membership_log( + "H3", + "membership redis hit after l1 miss", + {"prisma": False, "source": "redis"}, + ) + # #endregion + return cast(LiteLLM_TeamMembership | None, redis_membership) + + # #region agent log + _debug_team_membership_log("H1", "membership prisma fetch", {"prisma": True}) + # #endregion + return await _fetch_team_membership_from_db( + user_id=user_id, + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception: + verbose_proxy_logger.exception( + "Error getting team membership for user_id: %s, team_id: %s", + user_id, + team_id, + ) + return None + + async def get_team_membership( user_id: str, team_id: str, @@ -2155,54 +2314,43 @@ async def get_team_membership( Do a isolated check for team membership vs. doing a combined key + team + user + team-membership check, as key might come in frequently for different users/teams. Larger call will slowdown query time. This way we get to cache the constant (key/team/user info) and only update based on the changing value (team membership). """ - from litellm.proxy._types import LiteLLM_TeamMembership - - if prisma_client is None: - raise Exception("No db connected") - if user_id is None or team_id is None: return None _key: Final = team_membership_reservation_cache_key(user_id=user_id, team_id=team_id) - # check if in cache - cached: Final[object] = await user_api_key_cache.async_get_cache(key=_key) - if cached == NO_TEAM_MEMBERSHIP_SENTINEL: - return None - cached_membership_obj: Final = CacheCodec.deserialize(cached, model_type=LiteLLM_TeamMembership) - if cached_membership_obj is not None: - return cached_membership_obj + l1_cached: Final[object] = await user_api_key_cache.async_get_cache(key=_key, local_only=True) + l1_membership: Final = _membership_from_cached_payload(l1_cached) + if l1_membership is not _TEAM_MEMBERSHIP_CACHE_MISS: + # #region agent log + _debug_team_membership_log("H4", "membership l1 hit", {"prisma": False, "source": "l1"}) + # #endregion + return cast(LiteLLM_TeamMembership | None, l1_membership) - # else, check db - try: - response: Final = await _dictable_table(TeamMembershipRepository(prisma_client)).find_unique( - where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, - include={"litellm_budget_table": True}, + inflight: Final = _team_membership_inflight.get(_key) + if inflight is not None: + # #region agent log + _debug_team_membership_log("H5", "membership coalesced waiter", {"prisma": False, "coalesced": True}) + # #endregion + return await inflight + + if prisma_client is None: + raise Exception("No db connected") + + task: Final = asyncio.ensure_future( + _load_team_membership_on_cache_miss( + user_id=user_id, + team_id=team_id, + cache_key=_key, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, ) - - if response is None: - await user_api_key_cache.async_set_cache( - key=_key, - value=NO_TEAM_MEMBERSHIP_SENTINEL, - ttl=get_management_object_ttl(user_api_key_cache), - ) - return None - - _response: Final = LiteLLM_TeamMembership.model_validate(response.dict()) - await user_api_key_cache.async_set_cache( - key=_key, - value=_response, - model_type=LiteLLM_TeamMembership, - ) - - return _response - except Exception: - verbose_proxy_logger.exception( - "Error getting team membership for user_id: %s, team_id: %s", - user_id, - team_id, - ) - return None + ) + _team_membership_inflight[_key] = task + task.add_done_callback(lambda _t, k=_key: _team_membership_inflight.pop(k, None)) + return await task def model_in_access_group(model: str, team_models: list[str] | None, llm_router: Router | None) -> bool: @@ -4122,18 +4270,21 @@ async def _team_member_granted_models( prisma_client: PrismaClient, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, + team_membership: LiteLLM_TeamMembership | None = None, + team_membership_loaded: bool = False, ) -> Sequence[str]: """The member's own ``allowed_models`` scope; empty when the member is not narrowed below the team.""" if team_object is None or valid_token.user_id is None: return () - team_membership: Final = await get_team_membership( - user_id=valid_token.user_id, - team_id=team_object.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) + if not team_membership_loaded: + team_membership = await get_team_membership( + user_id=valid_token.user_id, + team_id=team_object.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) return () if team_membership is None else _member_allowed_models(team_membership) @@ -4169,6 +4320,8 @@ async def _granted_model_lists( prisma_client: PrismaClient, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, + team_membership: LiteLLM_TeamMembership | None = None, + team_membership_loaded: bool = False, ) -> tuple[Sequence[str], ...]: """One model allowlist per level that participates in authorizing the request.""" return ( @@ -4180,6 +4333,8 @@ async def _granted_model_lists( prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, + team_membership=team_membership, + team_membership_loaded=team_membership_loaded, ), project_object.models if project_object is not None else (), await _org_granted_models( @@ -4274,6 +4429,8 @@ async def collect_matched_model_access_groups( prisma_client: PrismaClient | None, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, + team_membership: LiteLLM_TeamMembership | None = None, + team_membership_loaded: bool = False, ) -> tuple[str, ...]: """ The budgeted model access groups that authorized this request, sorted and deduplicated. @@ -4319,6 +4476,8 @@ async def collect_matched_model_access_groups( prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, + team_membership=team_membership, + team_membership_loaded=team_membership_loaded, ) for granted_model in granted_models ) @@ -4334,6 +4493,8 @@ async def stamp_matched_model_access_groups( prisma_client: PrismaClient | None, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, + team_membership: LiteLLM_TeamMembership | None = None, + team_membership_loaded: bool = False, ) -> tuple[str, ...]: """Record the groups that authorized this request on its auth object, for the post-call spend writer and the reservation counters, and hand them back for the budget check.""" @@ -4350,6 +4511,8 @@ async def stamp_matched_model_access_groups( prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, + team_membership=team_membership, + team_membership_loaded=team_membership_loaded, ) except Exception as e: # noqa: BLE001 # fail-safe: attribution is spend telemetry, it must never break auth verbose_proxy_logger.debug("model access group attribution failed: %s", e) @@ -5152,6 +5315,8 @@ async def _check_team_member_budget( prisma_client: PrismaClient | None, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, + team_membership: LiteLLM_TeamMembership | None = None, + team_membership_loaded: bool = False, ): """Check if team member is over their max budget within the team.""" if ( @@ -5160,23 +5325,25 @@ async def _check_team_member_budget( and valid_token is not None and valid_token.user_id is not None ): - team_membership: Final = await get_team_membership( - user_id=valid_token.user_id, - team_id=team_object.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) + if not team_membership_loaded: + team_membership = await get_team_membership( + user_id=valid_token.user_id, + team_id=team_object.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + loaded_membership = team_membership # Per-member override wins; otherwise fall back to the team-level # default configured via team.metadata["team_member_budget_id"]. team_member_budget: float | None = None if ( - team_membership is not None - and team_membership.litellm_budget_table is not None - and team_membership.litellm_budget_table.max_budget is not None + loaded_membership is not None + and loaded_membership.litellm_budget_table is not None + and loaded_membership.litellm_budget_table.max_budget is not None ): - team_member_budget = team_membership.litellm_budget_table.max_budget + team_member_budget = loaded_membership.litellm_budget_table.max_budget else: default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id") if isinstance(default_budget_id, str): @@ -5195,7 +5362,7 @@ async def _check_team_member_budget( team_member_budget = default_budget.max_budget if team_member_budget is not None: - team_member_spend = (team_membership.spend if team_membership is not None else 0.0) or 0.0 + team_member_spend = (loaded_membership.spend if loaded_membership is not None else 0.0) or 0.0 # Read from cross-pod counter (Redis-first) if available from litellm.proxy.proxy_server import get_current_spend @@ -5224,6 +5391,8 @@ async def _check_team_member_model_access( prisma_client: Optional["PrismaClient"], user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, + team_membership: LiteLLM_TeamMembership | None = None, + team_membership_loaded: bool = False, ) -> None: """ Check if a team member's per-member model scope allows access to the requested model. @@ -5234,22 +5403,24 @@ async def _check_team_member_model_access( if valid_token.user_id is None or team_object.team_id is None: return - team_membership: Final = await get_team_membership( - user_id=valid_token.user_id, - team_id=team_object.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) + if not team_membership_loaded: + team_membership = await get_team_membership( + user_id=valid_token.user_id, + team_id=team_object.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + loaded_membership = team_membership if ( - team_membership is None - or team_membership.litellm_budget_table is None - or not team_membership.litellm_budget_table.allowed_models + loaded_membership is None + or loaded_membership.litellm_budget_table is None + or not loaded_membership.litellm_budget_table.allowed_models ): return # no per-member restriction — inherit team-level check - member_allowed_models: Final[list[str]] = team_membership.litellm_budget_table.allowed_models + member_allowed_models: Final[list[str]] = loaded_membership.litellm_budget_table.allowed_models try: _can_object_call_model( model=model, diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index ea2c925212d..755fe9efe2c 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -6454,6 +6454,133 @@ async def test_get_team_membership_reads_sentinel_as_no_membership_not_a_model() mock_prisma_client.db.litellm_teammembership.find_unique.assert_not_awaited() +@pytest.mark.asyncio +async def test_get_team_membership_coalesces_parallel_db_fetches(): + """Concurrent misses for the same member must share one Prisma round-trip.""" + from litellm.proxy.auth.auth_checks import get_team_membership + + started = asyncio.Event() + release = asyncio.Event() + membership_row = MagicMock() + membership_row.dict = lambda: {"user_id": "u-parallel", "team_id": "t-parallel", "spend": 1.0} + + async def _slow_find_unique(*args, **kwargs): + started.set() + await release.wait() + return membership_row + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(side_effect=_slow_find_unique) + cache = UserApiKeyCache() + + async def _load(): + return await get_team_membership( + user_id="u-parallel", + team_id="t-parallel", + prisma_client=mock_prisma_client, + user_api_key_cache=cache, + ) + + first = asyncio.create_task(_load()) + second = asyncio.create_task(_load()) + await started.wait() + await asyncio.sleep(0) + release.set() + results = await asyncio.gather(first, second) + + assert results[0] is not None and results[1] is not None + assert results[0].user_id == "u-parallel" + assert results[1].user_id == "u-parallel" + mock_prisma_client.db.litellm_teammembership.find_unique.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_get_team_membership_returns_before_redis_set_completes(): + """Auth must not wait on DualCache Redis SET; L1 is enough for the next lookup.""" + from litellm.proxy.auth.auth_checks import get_team_membership + + membership_row = MagicMock() + membership_row.dict = lambda: {"user_id": "u-redis", "team_id": "t-redis", "spend": 2.0} + + hang_redis_set = asyncio.Event() + + async def _hanging_redis_set(*args, **kwargs): + await hang_redis_set.wait() + + redis_cache = MagicMock() + redis_cache.async_get_cache = AsyncMock(return_value=None) + redis_cache.async_set_cache = AsyncMock(side_effect=_hanging_redis_set) + + cache = UserApiKeyCache(redis_cache=redis_cache) + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row) + + first = await asyncio.wait_for( + get_team_membership( + user_id="u-redis", + team_id="t-redis", + prisma_client=mock_prisma_client, + user_api_key_cache=cache, + ), + timeout=0.5, + ) + second = await get_team_membership( + user_id="u-redis", + team_id="t-redis", + prisma_client=mock_prisma_client, + user_api_key_cache=cache, + ) + + hang_redis_set.set() + await asyncio.sleep(0) + + assert first is not None and second is not None + assert first.spend == 2.0 + assert second.spend == 2.0 + mock_prisma_client.db.litellm_teammembership.find_unique.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_common_checks_calls_get_team_membership_once_per_request(): + """Model-access, attribution, and member-budget must reuse one membership load.""" + from fastapi import Request + + from litellm.proxy.auth.auth_checks import common_checks + + team = LiteLLM_TeamTable(team_id="t-once") + token = UserAPIKeyAuth(token="k-once", user_id="u-once", team_id="t-once", models=["gpt-4o-mini"]) + membership = MagicMock() + membership.litellm_budget_table = None + membership.spend = 0.0 + + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache()), + patch( + "litellm.proxy.auth.auth_checks.get_team_membership", + new_callable=AsyncMock, + return_value=membership, + ) as load_membership, + patch("litellm.proxy.proxy_server.get_current_spend", new_callable=AsyncMock, return_value=0.0), + ): + result = await common_checks( + request_body={"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hi"}]}, + team_object=team, + user_object=LiteLLM_UserTable(user_id="u-once"), + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/chat/completions", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=token, + request=MagicMock(spec=Request), + ) + + assert result is True + assert load_membership.await_count == 1 + + @pytest.mark.asyncio async def test_invalidate_team_member_spend_state_evicts_the_negative_cache_sentinel(): """ From 70ddc7e4928b19fe2d73cc97ff707e0919a8bceb Mon Sep 17 00:00:00 2001 From: Shivi Jain Date: Mon, 14 Sep 2026 21:28:08 +0530 Subject: [PATCH 045/112] fix(auth): load team membership once per request and skip prisma on an L1 hit common_checks was querying get_team_membership twice, and DualCache awaited Redis SET on the auth path, so LRU eviction plus a hung Redis write showed up as two postgres spans --- litellm/constants.py | 7 + litellm/proxy/auth/auth_checks.py | 163 +++++++++++------- .../proxy/auth/test_auth_checks.py | 160 +++++++++++++++++ 3 files changed, 265 insertions(+), 65 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 5751e6e46af..5c7a02d0743 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -2039,3 +2039,10 @@ BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY: Final = "batch_enqueued_token_limit" # Shared read-only empty mapping, for defaulting optional Mapping parameters without # constructing a fresh mutable dict at each call site. EMPTY_MAPPING: Final = MappingProxyType({}) + + +class TeamMembershipCacheMiss: + __slots__ = () + + +TEAM_MEMBERSHIP_CACHE_MISS: Final = TeamMembershipCacheMiss() diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index abcade6b1c8..0baa79fb95a 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -35,6 +35,8 @@ from litellm.constants import ( MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE, REGISTRY_ERROR_NEGATIVE_CACHE_TTL, TAG_REGISTRY_MAX_SIZE, + TEAM_MEMBERSHIP_CACHE_MISS, + TeamMembershipCacheMiss, ) from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider @@ -327,12 +329,31 @@ _safe_json_loads_obj: Final = _typed_json_loads(safe_json_loads) last_db_access_time: Final = LimitedSizeOrderedDict(max_size=100) db_cache_expiry: Final = DEFAULT_IN_MEMORY_TTL # refresh every 5s -_TEAM_MEMBERSHIP_CACHE_MISS: Final = object() -_team_membership_inflight: dict[str, asyncio.Task[LiteLLM_TeamMembership | None]] = {} +_TEAM_MEMBERSHIP_INFLIGHT_MAX: Final = 10000 +_team_membership_inflight: Final = LimitedSizeOrderedDict(max_size=_TEAM_MEMBERSHIP_INFLIGHT_MAX) +_team_membership_write_epoch: Final = LimitedSizeOrderedDict(max_size=_TEAM_MEMBERSHIP_INFLIGHT_MAX) all_routes: Final = LiteLLMRoutes.openai_routes.value + LiteLLMRoutes.management_routes.value +def _membership_write_epoch(key: str) -> int: + cached: Final[object] = _team_membership_write_epoch.get(key, 0) + return cached if isinstance(cached, int) else 0 + + +def _bump_membership_write_epoch(key: str) -> None: + _team_membership_write_epoch[key] = _membership_write_epoch(key) + 1 + + +def _membership_from_shared_load(result: object) -> LiteLLM_TeamMembership | None: + if result is None or isinstance(result, LiteLLM_TeamMembership): + return result + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail="Failed to load team membership", + ) + + def _log_budget_lookup_failure(entity: str, error: Exception) -> None: """ Log a warning when budget lookup fails; cache will not be populated. @@ -2165,38 +2186,39 @@ async def get_tag_object( return tag_objects.get(tag_name) -def _debug_team_membership_log(hypothesis_id: str, message: str, data: dict[str, object]) -> None: - # #region agent log - try: - import json as _json - - with open("/Users/shijain/genai-apps/genai-proxy/.cursor/debug-86534f.log", "a", encoding="utf-8") as _f: - _f.write( - _json.dumps( - { - "sessionId": "86534f", - "timestamp": int(time.time() * 1000), - "location": "auth_checks.py:get_team_membership", - "message": message, - "hypothesisId": hypothesis_id, - "data": data, - } - ) - + "\n" - ) - except Exception: - pass - # #endregion - - -def _membership_from_cached_payload(cached: object) -> LiteLLM_TeamMembership | None | object: - """Decode a DualCache payload. ``_TEAM_MEMBERSHIP_CACHE_MISS`` means try the next tier.""" +def _membership_from_cached_payload( + cached: object, +) -> LiteLLM_TeamMembership | None | TeamMembershipCacheMiss: if cached is None: - return _TEAM_MEMBERSHIP_CACHE_MISS + return TEAM_MEMBERSHIP_CACHE_MISS if cached == NO_TEAM_MEMBERSHIP_SENTINEL: return None cached_membership: Final = CacheCodec.deserialize(cached, model_type=LiteLLM_TeamMembership) - return cached_membership if cached_membership is not None else _TEAM_MEMBERSHIP_CACHE_MISS + return cached_membership if cached_membership is not None else TEAM_MEMBERSHIP_CACHE_MISS + + +async def _set_team_membership_cache_entry( + user_api_key_cache: UserApiKeyCache, + key: str, + value: object, + *, + local_only: bool, + model_type: type[LiteLLM_TeamMembership] | None, + ttl: float | None, +) -> None: + match (model_type is not None, ttl is not None): + case (False, False): + await user_api_key_cache.async_set_cache(key=key, value=value, local_only=local_only) + case (False, True): + await user_api_key_cache.async_set_cache(key=key, value=value, local_only=local_only, ttl=ttl) + case (True, False): + await user_api_key_cache.async_set_cache( + key=key, value=value, local_only=local_only, model_type=model_type + ) + case (True, True): + await user_api_key_cache.async_set_cache( + key=key, value=value, local_only=local_only, model_type=model_type, ttl=ttl + ) async def _populate_team_membership_cache( @@ -2207,17 +2229,33 @@ async def _populate_team_membership_cache( model_type: type[LiteLLM_TeamMembership] | None = None, ttl: float | None = None, ) -> None: - """Await in-memory write; replicate to Redis off the auth await path.""" - kwargs: dict[str, object] = {} - if model_type is not None: - kwargs["model_type"] = model_type - if ttl is not None: - kwargs["ttl"] = ttl - await user_api_key_cache.async_set_cache(key=key, value=value, local_only=True, **kwargs) + write_epoch: Final = _membership_write_epoch(key) + await _set_team_membership_cache_entry( + user_api_key_cache, + key, + value, + local_only=True, + model_type=model_type, + ttl=ttl, + ) + if _membership_write_epoch(key) != write_epoch: + await user_api_key_cache.async_delete_cache(key) + return async def _replicate_to_redis() -> None: try: - await user_api_key_cache.async_set_cache(key=key, value=value, **kwargs) + if _membership_write_epoch(key) != write_epoch: + return + await _set_team_membership_cache_entry( + user_api_key_cache, + key, + value, + local_only=False, + model_type=model_type, + ttl=ttl, + ) + if _membership_write_epoch(key) != write_epoch: + await user_api_key_cache.async_delete_cache(key) except Exception: return @@ -2271,19 +2309,9 @@ async def _load_team_membership_on_cache_miss( try: redis_cached: Final[object] = await user_api_key_cache.async_get_cache(key=cache_key) redis_membership: Final = _membership_from_cached_payload(redis_cached) - if redis_membership is not _TEAM_MEMBERSHIP_CACHE_MISS: - # #region agent log - _debug_team_membership_log( - "H3", - "membership redis hit after l1 miss", - {"prisma": False, "source": "redis"}, - ) - # #endregion - return cast(LiteLLM_TeamMembership | None, redis_membership) + if not isinstance(redis_membership, TeamMembershipCacheMiss): + return redis_membership - # #region agent log - _debug_team_membership_log("H1", "membership prisma fetch", {"prisma": True}) - # #endregion return await _fetch_team_membership_from_db( user_id=user_id, team_id=team_id, @@ -2292,13 +2320,18 @@ async def _load_team_membership_on_cache_miss( parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, ) - except Exception: + except HTTPException: + raise + except Exception as e: verbose_proxy_logger.exception( "Error getting team membership for user_id: %s, team_id: %s", user_id, team_id, ) - return None + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail="Failed to load team membership", + ) from e async def get_team_membership( @@ -2321,18 +2354,12 @@ async def get_team_membership( l1_cached: Final[object] = await user_api_key_cache.async_get_cache(key=_key, local_only=True) l1_membership: Final = _membership_from_cached_payload(l1_cached) - if l1_membership is not _TEAM_MEMBERSHIP_CACHE_MISS: - # #region agent log - _debug_team_membership_log("H4", "membership l1 hit", {"prisma": False, "source": "l1"}) - # #endregion - return cast(LiteLLM_TeamMembership | None, l1_membership) + if not isinstance(l1_membership, TeamMembershipCacheMiss): + return l1_membership - inflight: Final = _team_membership_inflight.get(_key) - if inflight is not None: - # #region agent log - _debug_team_membership_log("H5", "membership coalesced waiter", {"prisma": False, "coalesced": True}) - # #endregion - return await inflight + inflight: Final[object] = _team_membership_inflight.get(_key) + if isinstance(inflight, asyncio.Task): + return _membership_from_shared_load(await asyncio.shield(inflight)) if prisma_client is None: raise Exception("No db connected") @@ -2349,8 +2376,13 @@ async def get_team_membership( ) ) _team_membership_inflight[_key] = task - task.add_done_callback(lambda _t, k=_key: _team_membership_inflight.pop(k, None)) - return await task + + def _clear_inflight(_done: object) -> None: + if _team_membership_inflight.get(_key) is task: + _team_membership_inflight.pop(_key, None) + + task.add_done_callback(_clear_inflight) + return _membership_from_shared_load(await asyncio.shield(task)) def model_in_access_group(model: str, team_models: list[str] | None, llm_router: Router | None) -> bool: @@ -2861,6 +2893,7 @@ async def invalidate_team_member_spend_state( ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS, ) + _bump_membership_write_epoch(team_membership_reservation_cache_key(user_id=user_id, team_id=team_id)) await evict_and_broadcast( cache_keys=( team_membership_auth_cache_key(team_id=team_id, user_id=user_id), diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 755fe9efe2c..d16c75a4a3b 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -6581,6 +6581,166 @@ async def test_common_checks_calls_get_team_membership_once_per_request(): assert load_membership.await_count == 1 +@pytest.mark.asyncio +async def test_get_team_membership_db_error_raises_503_not_none(): + """A Prisma failure must fail closed as 503, not look like a missing membership row.""" + from fastapi import HTTPException + + from litellm.proxy.auth.auth_checks import get_team_membership + from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(side_effect=RuntimeError("db down")) + cache = UserApiKeyCache() + + with pytest.raises(HTTPException) as exc: + await get_team_membership( + user_id="u-fail", + team_id="t-fail", + prisma_client=mock_prisma_client, + user_api_key_cache=cache, + ) + + cached = await cache.async_get_cache( + key=team_membership_reservation_cache_key(user_id="u-fail", team_id="t-fail") + ) + assert exc.value.status_code == status.HTTP_503_SERVICE_UNAVAILABLE + assert cached is None + + +@pytest.mark.asyncio +async def test_common_checks_does_not_skip_member_limits_when_membership_lookup_fails(): + """common_checks must not mark membership loaded-absent after a lookup error.""" + from fastapi import HTTPException, Request + + from litellm.proxy.auth.auth_checks import common_checks + + team = LiteLLM_TeamTable(team_id="t-fail-closed") + token = UserAPIKeyAuth( + token="k-fail-closed", + user_id="u-fail-closed", + team_id="t-fail-closed", + models=["gpt-4o-mini"], + ) + lookup_error = HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail="Failed to load team membership", + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache()), + patch( + "litellm.proxy.auth.auth_checks.get_team_membership", + new_callable=AsyncMock, + side_effect=lookup_error, + ), + patch("litellm.proxy.proxy_server.get_current_spend", new_callable=AsyncMock, return_value=0.0), + ): + with pytest.raises(HTTPException) as exc: + await common_checks( + request_body={"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hi"}]}, + team_object=team, + user_object=LiteLLM_UserTable(user_id="u-fail-closed"), + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/chat/completions", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=token, + request=MagicMock(spec=Request), + ) + + assert exc.value.status_code == status.HTTP_503_SERVICE_UNAVAILABLE + + +@pytest.mark.asyncio +async def test_get_team_membership_waiter_cancel_does_not_cancel_shared_load(): + """Cancelling one coalesced waiter must not cancel the shared Prisma load.""" + from litellm.proxy.auth.auth_checks import get_team_membership + + started = asyncio.Event() + release = asyncio.Event() + membership_row = MagicMock() + membership_row.dict = lambda: {"user_id": "u-shield", "team_id": "t-shield", "spend": 1.0} + + async def _slow_find_unique(*args, **kwargs): + started.set() + await release.wait() + return membership_row + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(side_effect=_slow_find_unique) + cache = UserApiKeyCache() + + async def _load(): + return await get_team_membership( + user_id="u-shield", + team_id="t-shield", + prisma_client=mock_prisma_client, + user_api_key_cache=cache, + ) + + owner = asyncio.create_task(_load()) + await started.wait() + waiter = asyncio.create_task(_load()) + await asyncio.sleep(0) + waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await waiter + release.set() + result = await owner + + assert result is not None + assert result.user_id == "u-shield" + mock_prisma_client.db.litellm_teammembership.find_unique.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_stale_membership_redis_replicate_does_not_restore_after_invalidate(): + """A delayed DualCache Redis SET must not resurrect membership after invalidation.""" + from litellm.proxy.auth.auth_checks import get_team_membership, invalidate_team_member_spend_state + from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key + + membership_row = MagicMock() + membership_row.dict = lambda: {"user_id": "u-stale", "team_id": "t-stale", "spend": 9.0} + + hang_redis_set = asyncio.Event() + + async def _hanging_redis_set(*args, **kwargs): + await hang_redis_set.wait() + + redis_cache = MagicMock() + redis_cache.async_get_cache = AsyncMock(return_value=None) + redis_cache.async_set_cache = AsyncMock(side_effect=_hanging_redis_set) + redis_cache.async_delete_cache = AsyncMock() + redis_cache.delete_cache = MagicMock() + + cache = UserApiKeyCache(redis_cache=redis_cache) + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row) + + loaded = await get_team_membership( + user_id="u-stale", + team_id="t-stale", + prisma_client=mock_prisma_client, + user_api_key_cache=cache, + ) + cache_key = team_membership_reservation_cache_key(user_id="u-stale", team_id="t-stale") + await invalidate_team_member_spend_state(user_id="u-stale", team_id="t-stale", user_api_key_cache=cache) + after_invalidate = await cache.async_get_cache(key=cache_key, local_only=True) + + hang_redis_set.set() + await asyncio.sleep(0) + await asyncio.sleep(0) + after_replicate = await cache.async_get_cache(key=cache_key, local_only=True) + + assert loaded is not None + assert after_invalidate is None + assert after_replicate is None + + @pytest.mark.asyncio async def test_invalidate_team_member_spend_state_evicts_the_negative_cache_sentinel(): """ From 53ba8b986644b80fc95d39d11381d88eddb87309 Mon Sep 17 00:00:00 2001 From: Shivi Jain Date: Mon, 14 Sep 2026 22:05:17 +0530 Subject: [PATCH 046/112] fix(auth): load team membership once per request and skip prisma on an L1 hit common_checks was querying get_team_membership twice, and DualCache awaited Redis SET on the auth path, so LRU eviction plus a hung Redis write showed up as two postgres spans --- litellm/proxy/auth/auth_checks.py | 73 +++++++++++-------- litellm/proxy/auth/auth_object_prefetch.py | 5 +- .../proxy/auth/test_auth_checks.py | 34 ++++++--- .../proxy/auth/test_auth_object_prefetch.py | 4 +- .../auth/test_model_access_group_budgets.py | 34 ++++++--- 5 files changed, 93 insertions(+), 57 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 0baa79fb95a..89b061bdefe 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2197,30 +2197,51 @@ def _membership_from_cached_payload( return cached_membership if cached_membership is not None else TEAM_MEMBERSHIP_CACHE_MISS -async def _set_team_membership_cache_entry( +async def _set_team_membership_l1( user_api_key_cache: UserApiKeyCache, key: str, value: object, *, - local_only: bool, model_type: type[LiteLLM_TeamMembership] | None, ttl: float | None, ) -> None: match (model_type is not None, ttl is not None): case (False, False): - await user_api_key_cache.async_set_cache(key=key, value=value, local_only=local_only) + await user_api_key_cache.async_set_cache(key=key, value=value, local_only=True) case (False, True): - await user_api_key_cache.async_set_cache(key=key, value=value, local_only=local_only, ttl=ttl) + await user_api_key_cache.async_set_cache(key=key, value=value, local_only=True, ttl=ttl) case (True, False): - await user_api_key_cache.async_set_cache( - key=key, value=value, local_only=local_only, model_type=model_type - ) + await user_api_key_cache.async_set_cache(key=key, value=value, local_only=True, model_type=model_type) case (True, True): await user_api_key_cache.async_set_cache( - key=key, value=value, local_only=local_only, model_type=model_type, ttl=ttl + key=key, value=value, local_only=True, model_type=model_type, ttl=ttl ) +async def _replicate_team_membership_to_redis( + user_api_key_cache: UserApiKeyCache, + key: str, + value: object, + *, + model_type: type[LiteLLM_TeamMembership] | None, + ttl: float | None, + write_epoch: int, +) -> None: + redis_cache: Final = user_api_key_cache.redis_cache + if redis_cache is None or _membership_write_epoch(key) != write_epoch: + return + payload: Final[object] = CacheCodec.serialize(value, model_type=model_type) + try: + if ttl is None: + await redis_cache.async_set_cache(key, payload) + else: + await redis_cache.async_set_cache(key, payload, ttl=ttl) + if _membership_write_epoch(key) != write_epoch: + await redis_cache.async_delete_cache(key) + except Exception: + return + + async def _populate_team_membership_cache( user_api_key_cache: UserApiKeyCache, key: str, @@ -2230,36 +2251,27 @@ async def _populate_team_membership_cache( ttl: float | None = None, ) -> None: write_epoch: Final = _membership_write_epoch(key) - await _set_team_membership_cache_entry( + await _set_team_membership_l1( user_api_key_cache, key, value, - local_only=True, model_type=model_type, ttl=ttl, ) if _membership_write_epoch(key) != write_epoch: - await user_api_key_cache.async_delete_cache(key) + user_api_key_cache.in_memory_cache_for(key).delete_cache(key) return - async def _replicate_to_redis() -> None: - try: - if _membership_write_epoch(key) != write_epoch: - return - await _set_team_membership_cache_entry( - user_api_key_cache, - key, - value, - local_only=False, - model_type=model_type, - ttl=ttl, - ) - if _membership_write_epoch(key) != write_epoch: - await user_api_key_cache.async_delete_cache(key) - except Exception: - return - - asyncio.create_task(_replicate_to_redis()) + asyncio.create_task( + _replicate_team_membership_to_redis( + user_api_key_cache, + key, + value, + model_type=model_type, + ttl=ttl, + write_epoch=write_epoch, + ) + ) @log_db_metrics @@ -2363,6 +2375,9 @@ async def get_team_membership( if prisma_client is None: raise Exception("No db connected") + prisma: Final[object] = prisma_client + if isinstance(prisma, str): + return None task: Final = asyncio.ensure_future( _load_team_membership_on_cache_miss( diff --git a/litellm/proxy/auth/auth_object_prefetch.py b/litellm/proxy/auth/auth_object_prefetch.py index 52e26e885c9..5efc45d3fce 100644 --- a/litellm/proxy/auth/auth_object_prefetch.py +++ b/litellm/proxy/auth/auth_object_prefetch.py @@ -14,7 +14,6 @@ from pydantic import BaseModel, TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger from litellm.caching.redis_cache import RedisCache -from litellm.constants import DEFAULT_IN_MEMORY_TTL from litellm.models.organization import LiteLLM_OrganizationTable from litellm.models.team import LiteLLM_TeamTableCachedObj from litellm.models.team_membership import LiteLLM_TeamMembership @@ -191,13 +190,13 @@ def _iter_entries(refs: AuthObjectRefs, management_ttl: float) -> Iterator[_Cach ) if refs.organization_id is not None: yield _CacheEntry( - f"org_id:{refs.organization_id}", "organization_row", LiteLLM_OrganizationTable, DEFAULT_IN_MEMORY_TTL + f"org_id:{refs.organization_id}", "organization_row", LiteLLM_OrganizationTable, management_ttl ) yield _CacheEntry( f"org_id:{refs.organization_id}:with_budget", "organization_row", LiteLLM_OrganizationTable, - DEFAULT_IN_MEMORY_TTL, + management_ttl, ) if refs.project_id is not None: yield _CacheEntry(f"project_id:{refs.project_id}", "project_row", LiteLLM_ProjectTableCachedObj, management_ttl) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index d16c75a4a3b..4fdb49e47ef 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -5527,7 +5527,9 @@ async def _run_internal_user_budget_alert( with ( patch("litellm.proxy.proxy_server.prisma_client", None), # test-quality-ok: common_checks has no database seam - patch("litellm.proxy.proxy_server.get_current_spend", _get_spend), # test-quality-ok: common_checks imports it locally + patch( + "litellm.proxy.proxy_server.get_current_spend", _get_spend + ), # test-quality-ok: common_checks imports it locally patch.object(slack_alerting, "send_alert", send_alert), ): error: Final = await _check_for_error() @@ -6419,9 +6421,7 @@ async def test_get_team_membership_negative_caches_a_missing_row(): assert first is None assert second is None mock_prisma_client.db.litellm_teammembership.find_unique.assert_awaited_once() - cached = await cache.async_get_cache( - key=team_membership_reservation_cache_key(user_id="u-1", team_id="t-1") - ) + cached = await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-1", team_id="t-1")) assert cached == NO_TEAM_MEMBERSHIP_SENTINEL @@ -6601,13 +6601,25 @@ async def test_get_team_membership_db_error_raises_503_not_none(): user_api_key_cache=cache, ) - cached = await cache.async_get_cache( - key=team_membership_reservation_cache_key(user_id="u-fail", team_id="t-fail") - ) + cached = await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-fail", team_id="t-fail")) assert exc.value.status_code == status.HTTP_503_SERVICE_UNAVAILABLE assert cached is None +@pytest.mark.asyncio +async def test_get_team_membership_string_prisma_client_returns_none(): + """Unit tests stub prisma_client as a string; that is not a lookup failure and must not 503.""" + from litellm.proxy.auth.auth_checks import get_team_membership + + result = await get_team_membership( + user_id="u-str", + team_id="t-str", + prisma_client="hello-world", + user_api_key_cache=UserApiKeyCache(), + ) + assert result is None + + @pytest.mark.asyncio async def test_common_checks_does_not_skip_member_limits_when_membership_lookup_fails(): """common_checks must not mark membership loaded-absent after a lookup error.""" @@ -6699,7 +6711,7 @@ async def test_get_team_membership_waiter_cancel_does_not_cancel_shared_load(): @pytest.mark.asyncio async def test_stale_membership_redis_replicate_does_not_restore_after_invalidate(): - """A delayed DualCache Redis SET must not resurrect membership after invalidation.""" + """A delayed Redis SET must not rewrite L1 or leave Redis holding membership after invalidation.""" from litellm.proxy.auth.auth_checks import get_team_membership, invalidate_team_member_spend_state from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key @@ -6739,6 +6751,7 @@ async def test_stale_membership_redis_replicate_does_not_restore_after_invalidat assert loaded is not None assert after_invalidate is None assert after_replicate is None + redis_cache.async_delete_cache.assert_awaited() @pytest.mark.asyncio @@ -6764,10 +6777,7 @@ async def test_invalidate_team_member_spend_state_evicts_the_negative_cache_sent assert before is None await invalidate_team_member_spend_state(user_id="u-1", team_id="t-1", user_api_key_cache=cache) - assert ( - await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-1", team_id="t-1")) - is None - ) + assert await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-1", team_id="t-1")) is None after = await get_team_membership( user_id="u-1", team_id="t-1", prisma_client=mock_prisma_client, user_api_key_cache=cache diff --git a/tests/test_litellm/proxy/auth/test_auth_object_prefetch.py b/tests/test_litellm/proxy/auth/test_auth_object_prefetch.py index 0fd0dda3017..58761e270a0 100644 --- a/tests/test_litellm/proxy/auth/test_auth_object_prefetch.py +++ b/tests/test_litellm/proxy/auth/test_auth_object_prefetch.py @@ -181,8 +181,8 @@ async def test_cold_regime_is_one_mget_one_query_and_the_getters_never_touch_io_ assert sets == sorted( [ f"SET {TEAM_ID}_{USER_ID} ttl=5", - f"SET org_id:{ORG_ID} ttl=5", - f"SET org_id:{ORG_ID}:with_budget ttl=5", + f"SET org_id:{ORG_ID} ttl=60", + f"SET org_id:{ORG_ID}:with_budget ttl=60", f"SET {USER_ID} ttl=60", f"SET team_id:{TEAM_ID} ttl=60", f"SET team_membership:{USER_ID}:{TEAM_ID} ttl=None", diff --git a/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py b/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py index fb82d8708fd..a4206a0ac92 100644 --- a/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py +++ b/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py @@ -31,6 +31,7 @@ from litellm.proxy.auth.auth_checks import ( ) from litellm.proxy.common_utils.reset_budget_job import _model_access_group_counter_key from litellm.proxy.common_utils.user_api_key_cache import ( + NO_TEAM_MEMBERSHIP_SENTINEL, UserApiKeyCache, model_access_group_registry_cache_key, model_access_group_spend_counter_key, @@ -95,6 +96,11 @@ async def _cache( ), model_type=LiteLLM_TeamMembership, ) + else: + await cache.async_set_cache( + key=team_membership_reservation_cache_key(user_id=USER_ID, team_id=TEAM_ID), + value=NO_TEAM_MEMBERSHIP_SENTINEL, + ) if org_models: await cache.async_set_cache( key=f"org_id:{ORG_ID}", @@ -314,9 +320,7 @@ class _RecordingPrismaClient: def __init__(self, *rows: _MagBudgetRow) -> None: self.rows = {row.access_group_name: row for row in rows} self.batches: list[list[str]] = [] - self.db = SimpleNamespace( - litellm_modelaccessgroupbudgettable=SimpleNamespace(find_many=self._find_many) - ) + self.db = SimpleNamespace(litellm_modelaccessgroupbudgettable=SimpleNamespace(find_many=self._find_many)) async def _find_many(self, **kwargs): requested = list(kwargs["where"]["access_group_name"]["in"]) @@ -345,7 +349,9 @@ async def _enforce( read, seen = _spend_reader(spend_by_counter_key or {}) # The check takes its client and cache as arguments, injected just below. get_current_spend is the # one collaborator it reaches by a lazy `from litellm.proxy.proxy_server import`, with no parameter. - with patch("litellm.proxy.proxy_server.get_current_spend", read): # test-quality-ok: get_current_spend is lazily imported inside _model_access_group_max_budget_check and has no injection point + with patch( + "litellm.proxy.proxy_server.get_current_spend", read + ): # test-quality-ok: get_current_spend is lazily imported inside _model_access_group_max_budget_check and has no injection point await _model_access_group_max_budget_check( matched_model_access_groups=matched, prisma_client=prisma_client if prisma_client is not None else _RecordingPrismaClient(*rows), @@ -491,9 +497,7 @@ async def test_a_second_request_serves_the_budget_row_from_cache(): async def test_a_database_error_does_not_block_the_request(): class _FailingPrismaClient: def __init__(self) -> None: - self.db = SimpleNamespace( - litellm_modelaccessgroupbudgettable=SimpleNamespace(find_many=self._boom) - ) + self.db = SimpleNamespace(litellm_modelaccessgroupbudgettable=SimpleNamespace(find_many=self._boom)) async def _boom(self, **kwargs): raise RuntimeError("database unavailable") @@ -509,9 +513,15 @@ async def _common_checks_with_over_budget_group(*, skip_budget_checks: bool) -> with ( # common_checks resolves all three off the proxy_server module at call time; its signature # has no client, cache or spend-reader parameter to pass them through instead. - patch("litellm.proxy.proxy_server.prisma_client", prisma_client), # test-quality-ok: common_checks lazily imports prisma_client from proxy_server and takes no client parameter - patch("litellm.proxy.proxy_server.user_api_key_cache", cache), # test-quality-ok: common_checks lazily imports user_api_key_cache from proxy_server and takes no cache parameter - patch("litellm.proxy.proxy_server.get_current_spend", read), # test-quality-ok: get_current_spend is lazily imported inside the budget check and has no injection point + patch( + "litellm.proxy.proxy_server.prisma_client", prisma_client + ), # test-quality-ok: common_checks lazily imports prisma_client from proxy_server and takes no client parameter + patch( + "litellm.proxy.proxy_server.user_api_key_cache", cache + ), # test-quality-ok: common_checks lazily imports user_api_key_cache from proxy_server and takes no cache parameter + patch( + "litellm.proxy.proxy_server.get_current_spend", read + ), # test-quality-ok: get_current_spend is lazily imported inside the budget check and has no injection point ): return await common_checks( request_body={"model": "gpt-4o", "messages": []}, @@ -524,7 +534,9 @@ async def _common_checks_with_over_budget_group(*, skip_budget_checks: bool) -> llm_router=Router(model_list=MODEL_LIST), proxy_logging_obj=ProxyLogging(user_api_key_cache=cache), valid_token=UserAPIKeyAuth(api_key="hashed", models=["tier-a"], user_id=USER_ID), - request=SimpleNamespace(method="POST", headers={}, query_params={}, url=SimpleNamespace(path="/v1/chat/completions")), + request=SimpleNamespace( + method="POST", headers={}, query_params={}, url=SimpleNamespace(path="/v1/chat/completions") + ), skip_budget_checks=skip_budget_checks, ) From abaa2f8b81c4b73bc57a169fb2a85f6a535c0984 Mon Sep 17 00:00:00 2001 From: Shivi Jain Date: Mon, 14 Sep 2026 22:12:17 +0530 Subject: [PATCH 047/112] fix(auth): put TQ008 suppressions on the patch call lines The test-quality gate attributes the comment to the `patch(` line, so reasons on the closing paren did not count and lint failed after format started passing. Co-authored-by: Cursor --- .../proxy/auth/test_auth_checks.py | 32 +++++++++++++------ .../auth/test_model_access_group_budgets.py | 18 +++++------ 2 files changed, 30 insertions(+), 20 deletions(-) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 4fdb49e47ef..29f6c3861f6 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -5527,9 +5527,9 @@ async def _run_internal_user_budget_alert( with ( patch("litellm.proxy.proxy_server.prisma_client", None), # test-quality-ok: common_checks has no database seam - patch( + patch( # test-quality-ok: common_checks imports get_current_spend locally "litellm.proxy.proxy_server.get_current_spend", _get_spend - ), # test-quality-ok: common_checks imports it locally + ), patch.object(slack_alerting, "send_alert", send_alert), ): error: Final = await _check_for_error() @@ -6554,14 +6554,20 @@ async def test_common_checks_calls_get_team_membership_once_per_request(): membership.spend = 0.0 with ( - patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), - patch("litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache()), - patch( + patch( # test-quality-ok: common_checks imports prisma_client from proxy_server + "litellm.proxy.proxy_server.prisma_client", MagicMock() + ), + patch( # test-quality-ok: common_checks imports user_api_key_cache from proxy_server + "litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache() + ), + patch( # test-quality-ok: counts membership loads; common_checks has no membership seam "litellm.proxy.auth.auth_checks.get_team_membership", new_callable=AsyncMock, return_value=membership, ) as load_membership, - patch("litellm.proxy.proxy_server.get_current_spend", new_callable=AsyncMock, return_value=0.0), + patch( # test-quality-ok: common_checks imports get_current_spend locally + "litellm.proxy.proxy_server.get_current_spend", new_callable=AsyncMock, return_value=0.0 + ), ): result = await common_checks( request_body={"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hi"}]}, @@ -6640,14 +6646,20 @@ async def test_common_checks_does_not_skip_member_limits_when_membership_lookup_ ) with ( - patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), - patch("litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache()), - patch( + patch( # test-quality-ok: common_checks imports prisma_client from proxy_server + "litellm.proxy.proxy_server.prisma_client", MagicMock() + ), + patch( # test-quality-ok: common_checks imports user_api_key_cache from proxy_server + "litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache() + ), + patch( # test-quality-ok: injects membership lookup failure; common_checks has no seam "litellm.proxy.auth.auth_checks.get_team_membership", new_callable=AsyncMock, side_effect=lookup_error, ), - patch("litellm.proxy.proxy_server.get_current_spend", new_callable=AsyncMock, return_value=0.0), + patch( # test-quality-ok: common_checks imports get_current_spend locally + "litellm.proxy.proxy_server.get_current_spend", new_callable=AsyncMock, return_value=0.0 + ), ): with pytest.raises(HTTPException) as exc: await common_checks( diff --git a/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py b/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py index a4206a0ac92..7506bd031d9 100644 --- a/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py +++ b/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py @@ -349,9 +349,9 @@ async def _enforce( read, seen = _spend_reader(spend_by_counter_key or {}) # The check takes its client and cache as arguments, injected just below. get_current_spend is the # one collaborator it reaches by a lazy `from litellm.proxy.proxy_server import`, with no parameter. - with patch( + with patch( # test-quality-ok: get_current_spend is lazily imported inside the budget check "litellm.proxy.proxy_server.get_current_spend", read - ): # test-quality-ok: get_current_spend is lazily imported inside _model_access_group_max_budget_check and has no injection point + ): await _model_access_group_max_budget_check( matched_model_access_groups=matched, prisma_client=prisma_client if prisma_client is not None else _RecordingPrismaClient(*rows), @@ -511,17 +511,15 @@ async def _common_checks_with_over_budget_group(*, skip_budget_checks: bool) -> read, _ = _spend_reader({MODEL_ACCESS_GROUP_COUNTER_KEY: 99.0}) with ( - # common_checks resolves all three off the proxy_server module at call time; its signature - # has no client, cache or spend-reader parameter to pass them through instead. - patch( + patch( # test-quality-ok: common_checks lazily imports prisma_client from proxy_server "litellm.proxy.proxy_server.prisma_client", prisma_client - ), # test-quality-ok: common_checks lazily imports prisma_client from proxy_server and takes no client parameter - patch( + ), + patch( # test-quality-ok: common_checks lazily imports user_api_key_cache from proxy_server "litellm.proxy.proxy_server.user_api_key_cache", cache - ), # test-quality-ok: common_checks lazily imports user_api_key_cache from proxy_server and takes no cache parameter - patch( + ), + patch( # test-quality-ok: get_current_spend is lazily imported inside the budget check "litellm.proxy.proxy_server.get_current_spend", read - ), # test-quality-ok: get_current_spend is lazily imported inside the budget check and has no injection point + ), ): return await common_checks( request_body={"model": "gpt-4o", "messages": []}, From 75b16df6fa38d0521ee0bc658fb847611f865e06 Mon Sep 17 00:00:00 2001 From: Shivi Jain Date: Mon, 14 Sep 2026 22:22:34 +0530 Subject: [PATCH 048/112] fix(auth): keep prefetched org entries on the 5s getter TTL Organization mutations do not evict those cache keys, so stretching prefetch to the management TTL would leave stale org grants in L1. Co-authored-by: Cursor --- litellm/proxy/auth/auth_object_prefetch.py | 5 +++-- tests/test_litellm/proxy/auth/test_auth_object_prefetch.py | 4 ++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/auth/auth_object_prefetch.py b/litellm/proxy/auth/auth_object_prefetch.py index 5efc45d3fce..52e26e885c9 100644 --- a/litellm/proxy/auth/auth_object_prefetch.py +++ b/litellm/proxy/auth/auth_object_prefetch.py @@ -14,6 +14,7 @@ from pydantic import BaseModel, TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger from litellm.caching.redis_cache import RedisCache +from litellm.constants import DEFAULT_IN_MEMORY_TTL from litellm.models.organization import LiteLLM_OrganizationTable from litellm.models.team import LiteLLM_TeamTableCachedObj from litellm.models.team_membership import LiteLLM_TeamMembership @@ -190,13 +191,13 @@ def _iter_entries(refs: AuthObjectRefs, management_ttl: float) -> Iterator[_Cach ) if refs.organization_id is not None: yield _CacheEntry( - f"org_id:{refs.organization_id}", "organization_row", LiteLLM_OrganizationTable, management_ttl + f"org_id:{refs.organization_id}", "organization_row", LiteLLM_OrganizationTable, DEFAULT_IN_MEMORY_TTL ) yield _CacheEntry( f"org_id:{refs.organization_id}:with_budget", "organization_row", LiteLLM_OrganizationTable, - management_ttl, + DEFAULT_IN_MEMORY_TTL, ) if refs.project_id is not None: yield _CacheEntry(f"project_id:{refs.project_id}", "project_row", LiteLLM_ProjectTableCachedObj, management_ttl) diff --git a/tests/test_litellm/proxy/auth/test_auth_object_prefetch.py b/tests/test_litellm/proxy/auth/test_auth_object_prefetch.py index 58761e270a0..0fd0dda3017 100644 --- a/tests/test_litellm/proxy/auth/test_auth_object_prefetch.py +++ b/tests/test_litellm/proxy/auth/test_auth_object_prefetch.py @@ -181,8 +181,8 @@ async def test_cold_regime_is_one_mget_one_query_and_the_getters_never_touch_io_ assert sets == sorted( [ f"SET {TEAM_ID}_{USER_ID} ttl=5", - f"SET org_id:{ORG_ID} ttl=60", - f"SET org_id:{ORG_ID}:with_budget ttl=60", + f"SET org_id:{ORG_ID} ttl=5", + f"SET org_id:{ORG_ID}:with_budget ttl=5", f"SET {USER_ID} ttl=60", f"SET team_id:{TEAM_ID} ttl=60", f"SET team_membership:{USER_ID}:{TEAM_ID} ttl=None", From f4e21430c09d5201673b8093031d80391221a05a Mon Sep 17 00:00:00 2001 From: Shivi Jain Date: Mon, 14 Sep 2026 22:32:11 +0530 Subject: [PATCH 049/112] test(auth): freeze prefetch cache clock so the org getter cannot miss on a slow runner Prefetch still writes org entries with the 5s getter TTL. This test only asserts the SQL join, and wall-clock expiry on CI turned that into a MagicMock await TypeError. Co-authored-by: Cursor --- tests/proxy_behavior/auth/test_auth_object_prefetch.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/proxy_behavior/auth/test_auth_object_prefetch.py b/tests/proxy_behavior/auth/test_auth_object_prefetch.py index 59e9585a296..cfa958500af 100644 --- a/tests/proxy_behavior/auth/test_auth_object_prefetch.py +++ b/tests/proxy_behavior/auth/test_auth_object_prefetch.py @@ -66,6 +66,7 @@ async def test_join_binds_the_membership_to_the_requested_team(prisma): cache = _frozen_cache() refs = AuthObjectRefs(user_id=user_id, team_id=team_a, membership_user_id=user_id, organization_id=org_id) await prefetch_auth_objects(refs=refs, user_api_key_cache=cache, prisma_client=prisma) + assert cache.in_memory_cache.get_cache(f"org_id:{org_id}") is not None dead_db = _dead_db() membership = await get_team_membership( From cc1e8eb20f68350850e0337a1b8bfcfe9828447d Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 14 Sep 2026 11:19:28 -0700 Subject: [PATCH 050/112] fix: capture reused worker threads in Python traces --- .../shared/tracing/profiler.py | 68 +++++++---- .../shared/tracing/test_profiler.py | 113 ++++++++++++++++-- 2 files changed, 150 insertions(+), 31 deletions(-) diff --git a/tests/rust-python-harness/shared/tracing/profiler.py b/tests/rust-python-harness/shared/tracing/profiler.py index abfb6a2425d..e72a36a5289 100644 --- a/tests/rust-python-harness/shared/tracing/profiler.py +++ b/tests/rust-python-harness/shared/tracing/profiler.py @@ -2,7 +2,8 @@ from __future__ import annotations import sys import threading -from collections.abc import Generator, Iterator, Mapping +import warnings +from collections.abc import Callable, Generator, Iterator, Mapping from contextlib import contextmanager from dataclasses import dataclass from functools import lru_cache @@ -32,6 +33,7 @@ class PythonProfiler: self._source_root: Final = str(source_root.resolve()) + "/" self._seen_frames: Final[set[FrameType]] = set() self._event_ids: Final[dict[FrameType, int]] = {} + self._lock: Final = threading.Lock() self.events: Final[list[FunctionTraceEvent]] = [] def __call__(self, frame: FrameType, event: str, _arg: object) -> None: @@ -40,14 +42,15 @@ class PythonProfiler: function_name: Final = self.function_name(frame) if function_name is None: return - event_id: Final = len(self.events) - parent_id: Final = next( - (self._event_ids[ancestor] for ancestor in _frame_ancestors(frame) if ancestor in self._event_ids), - None, - ) - self._seen_frames.add(frame) - self._event_ids[frame] = event_id - self.events.append(FunctionTraceEvent(id=event_id, parent_id=parent_id, function=function_name)) + with self._lock: + event_id: Final = len(self.events) + parent_id: Final = next( + (self._event_ids[ancestor] for ancestor in _frame_ancestors(frame) if ancestor in self._event_ids), + None, + ) + self._seen_frames.add(frame) + self._event_ids[frame] = event_id + self.events.append(FunctionTraceEvent(id=event_id, parent_id=parent_id, function=function_name)) def function_name(self, frame: FrameType) -> str | None: code: Final = frame.f_code @@ -137,21 +140,51 @@ def _frame_ancestors(frame: FrameType) -> Generator[FrameType]: @contextmanager -def profile_python(source_root: Path, *, threads: bool = False) -> Generator[PythonProfiler]: - profiler: Final = PythonProfiler(source_root) +def _installed_profiler(profiler: Callable[[FrameType, str, object], None], *, threads: bool) -> Generator[None]: + if threads and sys.version_info >= (3, 12): + tool_id: Final = next((slot for slot in (2, 3, 4, 0, 1, 5) if sys.monitoring.get_tool(slot) is None), None) + if tool_id is None: + raise RuntimeError("no sys.monitoring tool ID is available for Python trace collection") + + def started(_code: CodeType, _offset: int) -> None: + profiler(sys._getframe(1), "call", None) + + sys.monitoring.use_tool_id(tool_id, "litellm-python-trace") + try: + sys.monitoring.register_callback(tool_id, sys.monitoring.events.PY_START, started) + sys.monitoring.set_events(tool_id, sys.monitoring.events.PY_START) + yield + finally: + sys.monitoring.set_events(tool_id, 0) + sys.monitoring.register_callback(tool_id, sys.monitoring.events.PY_START, None) + sys.monitoring.free_tool_id(tool_id) + return + if threads: + warnings.warn( + "Python <3.12 cannot trace existing worker threads; use Python 3.12+ for complete threaded traces", + RuntimeWarning, + stacklevel=3, + ) previous_thread: Final = threading.getprofile() if threads: threading.setprofile(profiler) previous: Final = sys.getprofile() sys.setprofile(profiler) try: - yield profiler + yield finally: sys.setprofile(previous) if threads: threading.setprofile(previous_thread) +@contextmanager +def profile_python(source_root: Path, *, threads: bool = False) -> Generator[PythonProfiler]: + profiler: Final = PythonProfiler(source_root) + with _installed_profiler(profiler, threads=threads): + yield profiler + + @contextmanager def profile_python_function_usage( source_root: Path, @@ -160,14 +193,5 @@ def profile_python_function_usage( threads: bool = False, ) -> Generator[PythonFunctionUsageProfiler]: profiler: Final = PythonFunctionUsageProfiler(source_root, functions) - previous_thread: Final = threading.getprofile() - if threads: - threading.setprofile(profiler) - previous: Final = sys.getprofile() - sys.setprofile(profiler) - try: + with _installed_profiler(profiler, threads=threads): yield profiler - finally: - sys.setprofile(previous) - if threads: - threading.setprofile(previous_thread) diff --git a/tests/rust-python-harness/shared/tracing/test_profiler.py b/tests/rust-python-harness/shared/tracing/test_profiler.py index 616e9c23e75..b9a07f249a3 100644 --- a/tests/rust-python-harness/shared/tracing/test_profiler.py +++ b/tests/rust-python-harness/shared/tracing/test_profiler.py @@ -4,9 +4,10 @@ import asyncio import sys import threading from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor from functools import wraps from pathlib import Path -from types import FunctionType +from types import FrameType, FunctionType from typing import Final, ParamSpec, TypeVar, cast import pytest @@ -41,11 +42,12 @@ def _events_named(profiler: PythonProfiler, name: str) -> tuple[FunctionTraceEve return tuple(event for event in profiler.events if event.function.endswith(name)) -def test_profiler_keeps_repeated_calls() -> None: +@pytest.mark.parametrize("threads", (False, True)) +def test_profiler_keeps_repeated_calls(threads: bool) -> None: def called() -> None: return None - with profile_python(Path(__file__).parent) as profiler: + with profile_python(Path(__file__).parent, threads=threads) as profiler: called() called() @@ -60,14 +62,15 @@ def test_profiler_qualifies_decorated_methods_by_class() -> None: assert _module_qualnames(__name__)[cast(FunctionType, Decorated.call.__wrapped__).__code__] == "Decorated.call" -def test_profiler_records_real_frame_ancestry() -> None: +@pytest.mark.parametrize("threads", (False, True)) +def test_profiler_records_real_frame_ancestry(threads: bool) -> None: def called() -> None: return None def outer() -> None: called() - with profile_python(Path(__file__).parent) as profiler: + with profile_python(Path(__file__).parent, threads=threads) as profiler: outer() outer_event, called_event = (event for event in profiler.events if event.function.endswith(("outer", "called"))) @@ -84,18 +87,20 @@ def test_profiler_restores_previous_profiler_after_failure() -> None: assert sys.getprofile() is previous -def test_profiler_does_not_count_coroutine_resumption_as_another_call() -> None: +@pytest.mark.parametrize("threads", (False, True)) +def test_profiler_does_not_count_coroutine_resumption_as_another_call(threads: bool) -> None: async def suspended() -> None: await asyncio.sleep(0) await asyncio.sleep(0) - with profile_python(Path(__file__).parent) as profiler: + with profile_python(Path(__file__).parent, threads=threads) as profiler: asyncio.run(suspended()) assert len(_events_named(profiler, "suspended")) == 1 -def test_profiler_preserves_parent_across_coroutine_suspension() -> None: +@pytest.mark.parametrize("threads", (False, True)) +def test_profiler_preserves_parent_across_coroutine_suspension(threads: bool) -> None: def called() -> None: return None @@ -103,7 +108,7 @@ def test_profiler_preserves_parent_across_coroutine_suspension() -> None: await asyncio.sleep(0) called() - with profile_python(Path(__file__).parent) as profiler: + with profile_python(Path(__file__).parent, threads=threads) as profiler: asyncio.run(suspended()) suspended_event: Final = _events_named(profiler, "suspended")[0] @@ -124,6 +129,96 @@ def test_profiler_captures_worker_threads_when_enabled() -> None: assert called_event.parent_id is None +@pytest.mark.skipif(sys.version_info < (3, 12), reason="existing worker capture requires sys.monitoring") +@pytest.mark.parametrize("prewarm", (False, True)) +def test_profiler_captures_reused_workers_without_leaking_between_sessions(prewarm: bool) -> None: + def called() -> None: + return None + + with ThreadPoolExecutor(max_workers=1) as executor: + if prewarm: + executor.submit(called).result(timeout=5) + with profile_python(Path(__file__).parent, threads=True) as first: + executor.submit(called).result(timeout=5) + executor.submit(called).result(timeout=5) + with profile_python(Path(__file__).parent, threads=True) as second: + executor.submit(called).result(timeout=5) + executor.submit(called).result(timeout=5) + + assert len(_events_named(first, "called")) == 1 + assert len(_events_named(second, "called")) == 1 + + +def test_profiler_restores_main_and_worker_hooks_after_failure() -> None: + previous: Final = sys.getprofile() + previous_thread: Final = threading.getprofile() + + with ThreadPoolExecutor(max_workers=1) as executor: + worker_previous: Final = executor.submit(sys.getprofile).result(timeout=5) + with pytest.raises(RuntimeError, match="stop"): + with profile_python(Path(__file__).parent, threads=True): + raise RuntimeError("stop") + assert executor.submit(sys.getprofile).result(timeout=5) is worker_previous + + assert sys.getprofile() is previous + assert threading.getprofile() is previous_thread + + +@pytest.mark.skipif(sys.version_info < (3, 12), reason="existing worker capture requires sys.monitoring") +def test_function_usage_profiler_captures_reused_workers() -> None: + def selected() -> None: + return None + + function: Final = f"{Path(__file__).name}:{selected.__code__.co_firstlineno} {selected.__qualname__}" + with ThreadPoolExecutor(max_workers=1) as executor: + executor.submit(selected).result(timeout=5) + with profile_python_function_usage(Path(__file__).parent, frozenset((function,)), threads=True) as profiler: + executor.submit(selected).result(timeout=5) + + assert profiler.called == {function} + + +@pytest.mark.skipif(sys.version_info < (3, 12), reason="independent thread hooks require sys.monitoring") +def test_threaded_profiler_preserves_custom_worker_hook_and_releases_monitoring_slot() -> None: + def worker_hook(_frame: FrameType, _event: str, _arg: object) -> None: + return None + + def fail_with_profile(executor: ThreadPoolExecutor) -> None: + with profile_python(Path(__file__).parent, threads=True): + assert executor.submit(sys.getprofile).result(timeout=5) is worker_hook + raise RuntimeError("stop") + + tools_before: Final = tuple(sys.monitoring.get_tool(slot) for slot in range(6)) + with ThreadPoolExecutor(max_workers=1, initializer=lambda: sys.setprofile(worker_hook)) as executor: + assert executor.submit(sys.getprofile).result(timeout=5) is worker_hook + with pytest.raises(RuntimeError, match="stop"): + fail_with_profile(executor) + assert executor.submit(sys.getprofile).result(timeout=5) is worker_hook + + assert tuple(sys.monitoring.get_tool(slot) for slot in range(6)) == tools_before + + +@pytest.mark.skipif(sys.version_info < (3, 12), reason="existing worker capture requires sys.monitoring") +def test_threaded_profiler_keeps_concurrent_event_ids_and_parent_links() -> None: + def child() -> None: + return None + + def parent() -> None: + child() + + with ThreadPoolExecutor(max_workers=4) as executor: + with profile_python(Path(__file__).parent, threads=True) as profiler: + futures: Final = tuple(executor.submit(parent) for _ in range(200)) + for future in futures: + future.result(timeout=5) + + parent_ids: Final = frozenset(event.id for event in _events_named(profiler, "parent")) + children: Final = _events_named(profiler, "child") + assert len(parent_ids) == len(children) == 200 + assert frozenset(event.parent_id for event in children) == parent_ids + assert tuple(event.id for event in profiler.events) == tuple(range(len(profiler.events))) + + def test_function_usage_profiler_records_only_selected_functions() -> None: def selected() -> None: return None From ed653ce9bf437e3e609b6fe3d1e9d2b3129aa417 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 14 Sep 2026 11:33:10 -0700 Subject: [PATCH 051/112] wip --- tests/rust-python-harness/AGENTS.md | 2 +- tests/rust-python-harness/cli/__init__.py | 15 +- tests/rust-python-harness/cli/catalog.py | 19 +- tests/rust-python-harness/cli/commands.py | 6 +- tests/rust-python-harness/cli/test_cli.py | 21 +- tests/rust-python-harness/conftest.py | 4 +- .../shared/native_build.py | 3 +- .../shared/parity/fixtures/recording.py | 1 + .../shared/parity/fixtures/store.py | 2 + .../shared/parity/fixtures/test_pipeline.py | 3 +- .../shared/parity/fixtures/test_recording.py | 1 + .../shared/reporting/models.py | 8 +- .../shared/reporting/strategy.py | 8 + .../shared/test_native_build.py | 16 +- .../shared/tracing/steps.py | 11 +- .../shared/tracing/test_steps.py | 4 +- .../e2e_parity/sdk/ocr/fixtures/reducto.py | 4 +- .../e2e_parity/sdk/ocr/test_sdk_parity.py | 5 +- .../strategies/trace_parity/AGENTS.md | 2 +- .../strategies/trace_parity/__init__.py | 34 ++- .../strategies/trace_parity/fixtures.py | 176 ++++++++++++ .../gateway/chat_completions/case.py | 65 +++++ .../trace_parity/gateway/execution.py | 61 ++-- .../trace_parity/gateway/messages/case.py | 48 ++-- .../trace_parity/gateway/responses/case.py | 63 ++++ .../strategies/trace_parity/models.py | 44 ++- .../strategies/trace_parity/reporting.py | 225 +++------------ .../strategies/trace_parity/runner.py | 90 +++--- .../trace_parity/sdk/chat_completions/case.py | 172 ++++++++--- .../strategies/trace_parity/sdk/execution.py | 81 +++--- .../trace_parity/sdk/messages/case.py | 269 ++++++++++++++++-- .../strategies/trace_parity/sdk/ocr/case.py | 91 ++++-- .../trace_parity/sdk/responses/case.py | 216 ++++++++++++++ .../sdk/test_core_scenario_matrix.py | 63 ++++ .../trace_parity/sdk/transcription/case.py | 13 +- .../strategies/trace_parity/test_reporting.py | 243 +++++----------- .../strategies/trace_parity/test_runner.py | 120 ++++++-- 37 files changed, 1554 insertions(+), 655 deletions(-) create mode 100644 tests/rust-python-harness/strategies/trace_parity/fixtures.py create mode 100644 tests/rust-python-harness/strategies/trace_parity/gateway/chat_completions/case.py create mode 100644 tests/rust-python-harness/strategies/trace_parity/gateway/responses/case.py create mode 100644 tests/rust-python-harness/strategies/trace_parity/sdk/responses/case.py create mode 100644 tests/rust-python-harness/strategies/trace_parity/sdk/test_core_scenario_matrix.py diff --git a/tests/rust-python-harness/AGENTS.md b/tests/rust-python-harness/AGENTS.md index 017668d4289..ec973035785 100644 --- a/tests/rust-python-harness/AGENTS.md +++ b/tests/rust-python-harness/AGENTS.md @@ -63,7 +63,7 @@ tests/rust-python-harness/ - Examples: `run e2e_parity --surface sdk --function ocr`, `run unit_tests_parity --function ocr --pytest-arg=-x`, or `run all --function ocr` - `cli/catalog.py` discovers strategies, validates their Python definitions, and orders them; `cli/__init__.py` builds the Click command tree; `cli/commands.py` runs selected cases - `e2e_parity/` compares SDK objects, exceptions, callbacks, and streams, or gateway HTTP responses -- `trace_parity/` compares mapped operations, call counts, and required execution ordering; before running it rebuilds the native bridge with the `trace-parity` feature whenever `litellm-rust` sources are newer than the installed extension (`shared/native_build.py`) +- `trace_parity/` prints filtered Python and Rust execution traces without comparing them; before running it rebuilds the native bridge with the `trace-parity` feature whenever `litellm-rust` sources are newer than the installed extension (`shared/native_build.py`) - E2E and trace strategies load their registered module cases and run surface-specific execution from their folders - `unit_tests_mapping/contracts.py` owns typed harness-side mapping contracts, per-function contracts live below `cases/`, and `mappings.py` exports the registry; live test discovery derives unmapped Python and Rust-only tests without an exhaustive manifest - `unit_tests_mapping/runner.py` validates confirmed mappings against the live Python and Rust inventories and attaches the derived status report diff --git a/tests/rust-python-harness/cli/__init__.py b/tests/rust-python-harness/cli/__init__.py index 13b995825dd..d2bfdc55b19 100644 --- a/tests/rust-python-harness/cli/__init__.py +++ b/tests/rust-python-harness/cli/__init__.py @@ -58,16 +58,29 @@ def _strategy_command(strategy: Strategy) -> click.Command: help=runner_argument.help, ) ) + for runner_option in strategy.definition.runner_options: + name: Final = runner_option.option.removeprefix("--").replace("-", "_") + params.append( + click.Option( + (runner_option.option, name), + type=click.Choice(runner_option.choices), + help=runner_option.help, + ) + ) def run_strategy( sdk_functions: tuple[str, ...], surface: str | None = None, runner_args: tuple[str, ...] = (), + **runner_options: str | None, ) -> int: selected_functions: Final = cast(frozenset[SdkFunction], frozenset(sdk_functions)) selected_surface: Final = cast(Surface | None, surface) cases: Final = select_cases((strategy,), selected_functions, selected_surface) - return run_command((strategy,), cases, runner_args) + option_args: Final = tuple( + f"--{name.replace('_', '-')}={value}" for name, value in runner_options.items() if value is not None + ) + return run_command((strategy,), cases, (*runner_args, *option_args)) return click.Command( strategy.id, diff --git a/tests/rust-python-harness/cli/catalog.py b/tests/rust-python-harness/cli/catalog.py index 03eb032d9c6..073873d0121 100644 --- a/tests/rust-python-harness/cli/catalog.py +++ b/tests/rust-python-harness/cli/catalog.py @@ -21,9 +21,7 @@ def _load_strategy_module(name: str, folder: Path, prefix: str | None) -> Module if prefix is not None: return importlib.import_module(f"{prefix}.{name}") module_name: Final = _synthetic_module_name(folder) - spec: Final = importlib.util.spec_from_file_location( - module_name, folder / "__init__.py" - ) + spec: Final = importlib.util.spec_from_file_location(module_name, folder / "__init__.py") if spec is None or spec.loader is None: raise ValueError(f"{folder}: cannot load strategy package") module: Final = importlib.util.module_from_spec(spec) @@ -59,9 +57,7 @@ def _load_strategy(name: str, folder: Path, prefix: str | None) -> Strategy: if duplicates: raise ValueError(f"{folder}: duplicate strategy cases: {duplicates}") expected: Final = frozenset( - (surface, function) - for surface in (definition.surfaces or (None,)) - for function in SDK_FUNCTIONS + (surface, function) for surface in (definition.surfaces or (None,)) for function in SDK_FUNCTIONS ) actual: Final = frozenset(keys) if actual != expected: @@ -73,8 +69,7 @@ def _load_strategy(name: str, folder: Path, prefix: str | None) -> Strategy: incompatible: Final = tuple( (case.surface, case.sdk_function) for case in definition.cases - if case.spec.disposition is CaseDisposition.RUNNABLE - and not isinstance(case.spec, definition.runnable_spec) + if case.spec.disposition is CaseDisposition.RUNNABLE and not isinstance(case.spec, definition.runnable_spec) ) if incompatible: raise ValueError(f"{folder}: runnable cases do not match {definition.runnable_spec.__name__}: {incompatible}") @@ -102,14 +97,10 @@ def _load_strategy(name: str, folder: Path, prefix: str | None) -> Strategy: def load_catalog(root: Path | None = None) -> tuple[Strategy, ...]: resolved: Final = STRATEGIES_ROOT if root is None else root prefix: Final = _STRATEGIES_PACKAGE.__name__ if resolved == STRATEGIES_ROOT else None - folders: Final = tuple( - info.name for info in pkgutil.iter_modules([str(resolved)]) if info.ispkg - ) + folders: Final = tuple(info.name for info in pkgutil.iter_modules([str(resolved)]) if info.ispkg) if not folders: raise ValueError(f"No strategy packages found below {resolved}") - strategies: Final = tuple( - _load_strategy(name, resolved / name, prefix) for name in sorted(folders) - ) + strategies: Final = tuple(_load_strategy(name, resolved / name, prefix) for name in sorted(folders)) ids: Final = [strategy.id for strategy in strategies] if len(set(ids)) != len(ids): raise ValueError(f"Duplicate strategy id in {resolved}") diff --git a/tests/rust-python-harness/cli/commands.py b/tests/rust-python-harness/cli/commands.py index f94c3277dc2..e51bbb5d966 100644 --- a/tests/rust-python-harness/cli/commands.py +++ b/tests/rust-python-harness/cli/commands.py @@ -21,8 +21,7 @@ def select_cases( case for strategy in strategies for case in strategy.cases - if (not sdk_functions or case.sdk_function in sdk_functions) - and (surface is None or case.surface == surface) + if (not sdk_functions or case.sdk_function in sdk_functions) and (surface is None or case.surface == surface) ) @@ -32,8 +31,7 @@ def run_command( runner_args: Sequence[str] = (), ) -> int: grouped: Final = { - strategy.id: tuple(case for case in cases if case.strategy_id == strategy.id) - for strategy in strategies + strategy.id: tuple(case for case in cases if case.strategy_id == strategy.id) for strategy in strategies } visible: Final = tuple(strategy for strategy in strategies if grouped[strategy.id]) runners: Final = tuple(replace(strategy, cases=grouped[strategy.id]) for strategy in visible) diff --git a/tests/rust-python-harness/cli/test_cli.py b/tests/rust-python-harness/cli/test_cli.py index 226c89843d0..5641aa8a539 100644 --- a/tests/rust-python-harness/cli/test_cli.py +++ b/tests/rust-python-harness/cli/test_cli.py @@ -242,7 +242,7 @@ def _assert_unavailable_cell(strategy: Strategy, case: HarnessCase, section_titl def test_every_unavailable_case_finishes_and_explains_itself() -> None: section_titles: Final = { "e2e_parity": "End-to-end parity outcomes", - "trace_parity": "trace comparisons", + "trace_parity": "traces", "unit_tests_mapping": "Python/Rust unit-test mappings", "unit_tests_parity": "Python backend parity outcomes", "unit_tests_rust": "Native Rust unit-test outcomes", @@ -359,6 +359,25 @@ def test_strategy_command_forwards_repeated_filters_and_runner_arguments( ] +def test_trace_command_forwards_engine_and_scenario(monkeypatch: pytest.MonkeyPatch) -> None: + cli: Final = importlib.import_module("tests.rust-python-harness.cli") + captured: list[tuple[str, ...]] = [] + + def capture_run( + strategies: Sequence[Strategy], + cases: Sequence[HarnessCase], + runner_args: Sequence[str] = (), + ) -> int: + del strategies, cases + captured.append(tuple(runner_args)) + return 0 + + monkeypatch.setattr(cli, "run_command", capture_run) + + assert main(["run", "trace_parity", "--scenario", "async-mistral", "--engine", "python"]) == 0 + assert captured == [("async-mistral", "--engine=python")] + + def test_omitted_surface_selects_every_strategy_surface(monkeypatch: pytest.MonkeyPatch) -> None: cli: Final = importlib.import_module("tests.rust-python-harness.cli") selected: list[str] = [] diff --git a/tests/rust-python-harness/conftest.py b/tests/rust-python-harness/conftest.py index d50d0fa4204..c487b25a9cc 100644 --- a/tests/rust-python-harness/conftest.py +++ b/tests/rust-python-harness/conftest.py @@ -19,9 +19,7 @@ def subprocess_test_environment(monkeypatch: pytest.MonkeyPatch) -> None: def cargo_project(tmp_path: Path) -> Callable[[str, str], Path]: def create(package: str, source: str) -> Path: manifest: Final = tmp_path / "Cargo.toml" - manifest.write_text( - f'[package]\nname = "{package}"\nversion = "0.1.0"\nedition = "2021"\n[workspace]\n' - ) + manifest.write_text(f'[package]\nname = "{package}"\nversion = "0.1.0"\nedition = "2021"\n[workspace]\n') (tmp_path / "src").mkdir() (tmp_path / "src/lib.rs").write_text(source) return manifest diff --git a/tests/rust-python-harness/shared/native_build.py b/tests/rust-python-harness/shared/native_build.py index 8693cf3bac2..b50def98e3a 100644 --- a/tests/rust-python-harness/shared/native_build.py +++ b/tests/rust-python-harness/shared/native_build.py @@ -85,7 +85,8 @@ def trace_bridge_error() -> str | None: def ensure_trace_bridge(repo_root: Path) -> str | None: native_path: Final = _native_module_path() native_mtime: Final = native_path.stat().st_mtime if native_path is not None and native_path.exists() else None - if needs_rebuild(native_mtime, _newest_source_mtime(repo_root)): + rebuild_required: Final = needs_rebuild(native_mtime, _newest_source_mtime(repo_root)) or trace_bridge_error() is not None + if rebuild_required: print(f"Rebuilding native Rust bridge ({BRIDGE_FEATURE} feature)...", flush=True) succeeded: Final output: Final diff --git a/tests/rust-python-harness/shared/parity/fixtures/recording.py b/tests/rust-python-harness/shared/parity/fixtures/recording.py index ee0d4b6de7b..cb6ed968b90 100644 --- a/tests/rust-python-harness/shared/parity/fixtures/recording.py +++ b/tests/rust-python-harness/shared/parity/fixtures/recording.py @@ -191,6 +191,7 @@ class _RecordingHandler(LocalHttpHandler): self.end_headers() self.wfile.write(body) + def _recording_provider(spec: UpstreamEndpoint) -> AbstractContextManager[_RecordingProvider]: return serve_in_thread(_RecordingProvider(spec)) diff --git a/tests/rust-python-harness/shared/parity/fixtures/store.py b/tests/rust-python-harness/shared/parity/fixtures/store.py index 270af2a7625..1145b5d7d27 100644 --- a/tests/rust-python-harness/shared/parity/fixtures/store.py +++ b/tests/rust-python-harness/shared/parity/fixtures/store.py @@ -15,6 +15,8 @@ from .cassette import deserialize_cassette, serialize_cassette from .recording import RecordedInteraction FIXTURE_SCHEMA_VERSION: Final = 1 + + class FixtureInput(Protocol): def canonical_input(self) -> dict[str, object]: ... diff --git a/tests/rust-python-harness/shared/parity/fixtures/test_pipeline.py b/tests/rust-python-harness/shared/parity/fixtures/test_pipeline.py index 4535ba05bf6..b162e4949d2 100644 --- a/tests/rust-python-harness/shared/parity/fixtures/test_pipeline.py +++ b/tests/rust-python-harness/shared/parity/fixtures/test_pipeline.py @@ -45,8 +45,8 @@ class _Upstream(LocalHttpServer): super().__init__(("127.0.0.1", 0), _UpstreamHandler) self.response_status: Final = status -class _UpstreamHandler(LocalHttpHandler): +class _UpstreamHandler(LocalHttpHandler): def do_POST(self) -> None: length: Final = int(self.headers.get("content-length") or "0") self.rfile.read(length) @@ -59,6 +59,7 @@ class _UpstreamHandler(LocalHttpHandler): self.end_headers() self.wfile.write(body) + def _upstream(status: int = 200) -> AbstractContextManager[_Upstream]: return serve_in_thread(_Upstream(status)) diff --git a/tests/rust-python-harness/shared/parity/fixtures/test_recording.py b/tests/rust-python-harness/shared/parity/fixtures/test_recording.py index 6181f18e89a..c1e7c2af8d8 100644 --- a/tests/rust-python-harness/shared/parity/fixtures/test_recording.py +++ b/tests/rust-python-harness/shared/parity/fixtures/test_recording.py @@ -238,6 +238,7 @@ class _ControlledUpstreamHandler(LocalHttpHandler): self.end_headers() self.wfile.write(body) + def _controlled_upstream( stream_chunks: tuple[bytes, ...] = _SSE_CHUNKS, ) -> AbstractContextManager[_ControlledUpstream]: diff --git a/tests/rust-python-harness/shared/reporting/models.py b/tests/rust-python-harness/shared/reporting/models.py index 1ebba6c9793..4a46a6d4cfe 100644 --- a/tests/rust-python-harness/shared/reporting/models.py +++ b/tests/rust-python-harness/shared/reporting/models.py @@ -180,15 +180,11 @@ class HarnessRun: @property def unique_checks(self) -> int: - return len( - {nodeid for result in self.results.values() for nodeid in result.collected} - ) + return len({nodeid for result in self.results.values() for nodeid in result.collected}) @property def completed_checks(self) -> int: - return len( - {nodeid for result in self.results.values() for nodeid in result.completed} - ) + return len({nodeid for result in self.results.values() for nodeid in result.completed}) @classmethod def from_cases(cls, cases: Iterable[HarnessCase]) -> HarnessRun: diff --git a/tests/rust-python-harness/shared/reporting/strategy.py b/tests/rust-python-harness/shared/reporting/strategy.py index 7e76f035e20..d8e9d9e5ba9 100644 --- a/tests/rust-python-harness/shared/reporting/strategy.py +++ b/tests/rust-python-harness/shared/reporting/strategy.py @@ -67,6 +67,13 @@ class RunnerArgumentDefinition: metavar: str = "ARG" +@dataclass(frozen=True, slots=True) +class RunnerOptionDefinition: + option: str + help: str + choices: tuple[str, ...] + + class StrategyRunner(Protocol): def __call__( self, @@ -90,3 +97,4 @@ class StrategyDefinition: render: StrategyRenderer surfaces: tuple[Surface, ...] = () runner_argument: RunnerArgumentDefinition | None = None + runner_options: tuple[RunnerOptionDefinition, ...] = () diff --git a/tests/rust-python-harness/shared/test_native_build.py b/tests/rust-python-harness/shared/test_native_build.py index f3e5aead846..982f787aefa 100644 --- a/tests/rust-python-harness/shared/test_native_build.py +++ b/tests/rust-python-harness/shared/test_native_build.py @@ -89,8 +89,8 @@ def test_ensure_trace_bridge_reports_failed_rebuild(tmp_path: Final, monkeypatch assert "boom" in message -def test_ensure_trace_bridge_flags_missing_trace_feature_without_rebuild( - tmp_path: Final, monkeypatch: pytest.MonkeyPatch +def test_ensure_trace_bridge_rebuilds_when_trace_feature_is_missing( + tmp_path: Final, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] ) -> None: native: Final = tmp_path / "_native.abi3.so" native.write_bytes(b"") @@ -107,10 +107,14 @@ def test_ensure_trace_bridge_flags_missing_trace_feature_without_rebuild( monkeypatch.setattr(native_build, "_native_module_path", lambda: native) monkeypatch.setattr(native_build, "_rebuild", fake_rebuild) - monkeypatch.setattr(native_build, "get_native_bridge", lambda: SimpleNamespace(_trace=None)) + monkeypatch.setattr( + native_build, + "get_native_bridge", + lambda: SimpleNamespace(_trace=object() if state.rebuilt else None), + ) message: Final = native_build.ensure_trace_bridge(tmp_path) - assert message is not None - assert "_trace" in message - assert state.rebuilt is False + assert message is None + assert state.rebuilt is True + assert "Rebuilding native Rust bridge" in capsys.readouterr().out diff --git a/tests/rust-python-harness/shared/tracing/steps.py b/tests/rust-python-harness/shared/tracing/steps.py index 2475f0fdcc5..152f753a536 100644 --- a/tests/rust-python-harness/shared/tracing/steps.py +++ b/tests/rust-python-harness/shared/tracing/steps.py @@ -181,12 +181,7 @@ class TraceDiff: @property def matches(self) -> bool: - return ( - not self.python_only - and not self.rust_only - and not self.missing_mappings - and self.shared_order_matches - ) + return not self.python_only and not self.rust_only and not self.missing_mappings and self.shared_order_matches def _missing_mappings( @@ -257,9 +252,7 @@ def trace_diff( rust_counts: Final = Counter(rust_spans) python_only_counts: Final = python_counts - rust_counts rust_only_counts: Final = rust_counts - python_counts - python_only: Final = tuple( - span for span, count in python_only_counts.items() for _ in range(count) - ) + python_only: Final = tuple(span for span, count in python_only_counts.items() for _ in range(count)) rust_only: Final = tuple(span for span, count in rust_only_counts.items() for _ in range(count)) first_difference: Final = _first_difference(python, rust, mappings, contract) return TraceDiff( diff --git a/tests/rust-python-harness/shared/tracing/test_steps.py b/tests/rust-python-harness/shared/tracing/test_steps.py index 2efc6a3c579..1b2f02a4ed7 100644 --- a/tests/rust-python-harness/shared/tracing/test_steps.py +++ b/tests/rust-python-harness/shared/tracing/test_steps.py @@ -146,9 +146,7 @@ def test_trace_diff_allows_reordered_concurrent_children() -> None: def test_trace_diff_prunes_declared_engine_only_nodes_but_requires_them() -> None: mappings: Final = (MAPPINGS[0], mapping(rust_span="rust_prepare")) python: Final = pipeline_projection("python", (event(0, "module.py:1 entry"),), mappings).steps - rust: Final = pipeline_projection( - "rust", (event(0, "route"), event(1, "rust_prepare", 0)), mappings - ).steps + rust: Final = pipeline_projection("rust", (event(0, "route"), event(1, "rust_prepare", 0)), mappings).steps assert trace_diff(python, rust, mappings).matches assert trace_diff(python, rust[:1], mappings).missing_mappings == ("rust_prepare",) diff --git a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/fixtures/reducto.py b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/fixtures/reducto.py index bcca8ac6d42..1b283d75ca5 100644 --- a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/fixtures/reducto.py +++ b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/fixtures/reducto.py @@ -264,9 +264,7 @@ def _formatting_strategy() -> SearchStrategy[ReductoFormatting]: ), st.sampled_from((False, True)).map(lambda value: {"add_page_markers": value}), st.sampled_from((False, True)).map(lambda value: {"merge_tables": value}), - st.sampled_from(REDUCTO_FORMATTING_INCLUDE_GROUPS) - .map(list) - .map(lambda value: {"include": value}), + st.sampled_from(REDUCTO_FORMATTING_INCLUDE_GROUPS).map(list).map(lambda value: {"include": value}), ) return values.map(ReductoFormatting.model_validate) diff --git a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py index 5c4abc78081..f14296fbf06 100644 --- a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py +++ b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py @@ -124,10 +124,7 @@ class RecordingCallback(CustomLogger): if isinstance(value, Mapping): if any(not isinstance(map_key, str) for map_key in value): raise TypeError("callback kwarg mappings must use string keys") - return { - map_key: self._normalized_kwargs(map_value, map_key) - for map_key, map_value in value.items() - } + return {map_key: self._normalized_kwargs(map_value, map_key) for map_key, map_value in value.items()} if isinstance(value, (list, tuple)): return [self._normalized_kwargs(item) for item in value] raise TypeError(f"unsupported callback kwarg type: {type(value)}") diff --git a/tests/rust-python-harness/strategies/trace_parity/AGENTS.md b/tests/rust-python-harness/strategies/trace_parity/AGENTS.md index bb7cb8c91d8..5db17974d07 100644 --- a/tests/rust-python-harness/strategies/trace_parity/AGENTS.md +++ b/tests/rust-python-harness/strategies/trace_parity/AGENTS.md @@ -1 +1 @@ -Maps Python profiler frames onto feature-gated Rust span names via an explicit per-case mapping (Rust span name is the identity) and compares steps, order, and nesting of both live traces against a replayed provider response. +Prints filtered Python profiler frames and feature-gated Rust spans from live traces against a replayed provider response. The two traces are independent and are not compared. diff --git a/tests/rust-python-harness/strategies/trace_parity/__init__.py b/tests/rust-python-harness/strategies/trace_parity/__init__.py index ec88b0169fa..710bdaa3d39 100644 --- a/tests/rust-python-harness/strategies/trace_parity/__init__.py +++ b/tests/rust-python-harness/strategies/trace_parity/__init__.py @@ -7,6 +7,7 @@ from ...shared.reporting.strategy import ( ModuleCaseSpec, NotImplementedCaseSpec, RunnerArgumentDefinition, + RunnerOptionDefinition, StrategyDefinition, ) from .reporting import render_trace_results @@ -26,13 +27,17 @@ CASES: Final[tuple[CaseDefinition, ...]] = ( ModuleCaseSpec( coverage=Coverage.PARTIAL, module="tests.rust-python-harness.strategies.trace_parity.sdk.messages.case", - note="Async only until anthropic_messages_handler supports sync calls.", + note="Success paths are async; sync tracing captures the currently unsupported behavior.", ), surface="sdk", ), CaseDefinition( "responses", - NotImplementedCaseSpec(reason="No Responses trace-parity case is registered."), + ModuleCaseSpec( + coverage=Coverage.PARTIAL, + module="tests.rust-python-harness.strategies.trace_parity.sdk.responses.case", + note="Core create paths: native, streaming, provider error, Azure override, and chat bridge.", + ), surface="sdk", ), CaseDefinition( @@ -70,13 +75,17 @@ CASES: Final[tuple[CaseDefinition, ...]] = ( ModuleCaseSpec( coverage=Coverage.PARTIAL, module="tests.rust-python-harness.strategies.trace_parity.gateway.messages.case", - note="Non-streaming success paths only.", + note="Anthropic/Azure provider routes plus a fully consumed downstream streaming path.", ), surface="gateway", ), CaseDefinition( "responses", - NotImplementedCaseSpec(reason="No gateway Responses trace-parity case is registered."), + ModuleCaseSpec( + coverage=Coverage.PARTIAL, + module="tests.rust-python-harness.strategies.trace_parity.gateway.responses.case", + note="Native OpenAI non-streaming and fully consumed downstream streaming paths.", + ), surface="gateway", ), CaseDefinition( @@ -86,7 +95,11 @@ CASES: Final[tuple[CaseDefinition, ...]] = ( ), CaseDefinition( "chat_completions", - NotImplementedCaseSpec(reason="No gateway chat trace-parity case is registered."), + ModuleCaseSpec( + coverage=Coverage.PARTIAL, + module="tests.rust-python-harness.strategies.trace_parity.gateway.chat_completions.case", + note="Anthropic non-streaming and fully consumed downstream streaming paths.", + ), surface="gateway", ), CaseDefinition( @@ -99,8 +112,8 @@ CASES: Final[tuple[CaseDefinition, ...]] = ( STRATEGY: Final = StrategyDefinition( id="trace_parity", order=20, - label="Trace parity", - description="Compare pipeline steps, order, and nesting between Python profiler frames and Rust spans via an explicit mapping.", + label="Traces", + description="Print Python profiler frames and Rust spans for representative pipeline scenarios.", directory=Path(__file__).parent, runnable_spec=ModuleCaseSpec, cases=CASES, @@ -112,4 +125,11 @@ STRATEGY: Final = StrategyDefinition( metavar="NAME", help="run only this named trace scenario; repeat to select more than one", ), + runner_options=( + RunnerOptionDefinition( + option="--engine", + choices=("python", "rust"), + help="show only this engine's trace; omit to print both engines", + ), + ), ) diff --git a/tests/rust-python-harness/strategies/trace_parity/fixtures.py b/tests/rust-python-harness/strategies/trace_parity/fixtures.py new file mode 100644 index 00000000000..84ac525efb6 --- /dev/null +++ b/tests/rust-python-harness/strategies/trace_parity/fixtures.py @@ -0,0 +1,176 @@ +from __future__ import annotations + +import base64 +import binascii +import json +import struct +from collections.abc import Iterable, Mapping +from typing import Final + +from ...shared.parity.recorded_http import ( + HttpHeader, + RecordedHttpResponse, + RecordedHttpStreamResponse, + RecordedStreamChunk, +) + +JSON_HEADERS: Final = (HttpHeader(name="content-type", value="application/json"),) +SSE_HEADERS: Final = (HttpHeader(name="content-type", value="text/event-stream"),) +AWS_EVENT_STREAM_HEADERS: Final = (HttpHeader(name="content-type", value="application/vnd.amazon.eventstream"),) + + +def json_response(body: Mapping[str, object] | bytes, *, status: int = 200) -> RecordedHttpResponse: + encoded: Final = body if isinstance(body, bytes) else json.dumps(body).encode() + return RecordedHttpResponse.from_bytes(status, JSON_HEADERS, encoded) + + +def sse_event(event: str, payload: Mapping[str, object]) -> bytes: + return f"event: {event}\ndata: {json.dumps(payload, separators=(',', ':'))}\n\n".encode() + + +def sse_response(events: Iterable[tuple[str, Mapping[str, object]]]) -> RecordedHttpStreamResponse: + return RecordedHttpStreamResponse( + kind="http_stream", + status_code=200, + headers=SSE_HEADERS, + chunks=tuple(RecordedStreamChunk.from_bytes(sse_event(event, payload)) for event, payload in events), + ) + + +def _aws_string_header(name: str, value: str) -> bytes: + name_bytes: Final = name.encode() + value_bytes: Final = value.encode() + return ( + struct.pack("!B", len(name_bytes)) + + name_bytes + + struct.pack("!B", 7) + + struct.pack("!H", len(value_bytes)) + + value_bytes + ) + + +def aws_event_stream_frame(payload: Mapping[str, object]) -> bytes: + event_payload: Final = json.dumps( + {"bytes": base64.b64encode(json.dumps(payload, separators=(",", ":")).encode()).decode()}, + separators=(",", ":"), + ).encode() + headers: Final = ( + _aws_string_header(":event-type", "chunk") + + _aws_string_header(":content-type", "application/json") + + _aws_string_header(":message-type", "event") + ) + total_length: Final = 12 + len(headers) + len(event_payload) + 4 + prelude: Final = struct.pack("!II", total_length, len(headers)) + prelude_crc: Final = binascii.crc32(prelude) & 0xFFFFFFFF + prelude_crc_bytes: Final = struct.pack("!I", prelude_crc) + message_crc: Final = binascii.crc32(prelude_crc_bytes + headers + event_payload, prelude_crc) & 0xFFFFFFFF + return prelude + prelude_crc_bytes + headers + event_payload + struct.pack("!I", message_crc) + + +def aws_event_stream_response( + events: Iterable[Mapping[str, object]], *, corrupt_last_frame: bool = False +) -> RecordedHttpStreamResponse: + frames: Final = [aws_event_stream_frame(event) for event in events] + if corrupt_last_frame: + corrupted: Final = bytearray(frames[-1]) + corrupted[-1] ^= 0xFF + frames[-1] = bytes(corrupted) + return RecordedHttpStreamResponse( + kind="http_stream", + status_code=200, + headers=AWS_EVENT_STREAM_HEADERS, + chunks=(RecordedStreamChunk.from_bytes(b"".join(frames)),), + ) + + +def anthropic_response_body(*, model: str = "claude-sonnet-5") -> dict[str, object]: + return { + "id": "msg_trace", + "type": "message", + "role": "assistant", + "model": model, + "content": [{"type": "text", "text": "hello"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 2, "output_tokens": 3}, + } + + +def anthropic_stream_events(*, model: str = "claude-sonnet-5") -> tuple[tuple[str, Mapping[str, object]], ...]: + return ( + ( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_trace", + "type": "message", + "role": "assistant", + "model": model, + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 2, "output_tokens": 0}, + }, + }, + ), + ( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + ( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hello"}}, + ), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 1}, + }, + ), + ("message_stop", {"type": "message_stop"}), + ) + + +def responses_body(*, model: str = "gpt-5", status: str = "completed") -> dict[str, object]: + return { + "id": "resp_trace", + "object": "response", + "created_at": 1_750_000_000, + "status": status, + "model": model, + "output": [ + { + "type": "message", + "id": "msg_trace", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "hello", "annotations": []}], + } + ], + "usage": {"input_tokens": 2, "output_tokens": 3, "total_tokens": 5}, + } + + +def responses_stream_events(*, model: str = "gpt-5") -> tuple[tuple[str, Mapping[str, object]], ...]: + response: Final = responses_body(model=model) + return ( + ( + "response.created", + {"type": "response.created", "response": {**response, "status": "in_progress", "output": []}}, + ), + ( + "response.output_text.delta", + { + "type": "response.output_text.delta", + "item_id": "msg_trace", + "output_index": 0, + "content_index": 0, + "delta": "hello", + }, + ), + ("response.completed", {"type": "response.completed", "response": response}), + ) diff --git a/tests/rust-python-harness/strategies/trace_parity/gateway/chat_completions/case.py b/tests/rust-python-harness/strategies/trace_parity/gateway/chat_completions/case.py new file mode 100644 index 00000000000..522d2d0259c --- /dev/null +++ b/tests/rust-python-harness/strategies/trace_parity/gateway/chat_completions/case.py @@ -0,0 +1,65 @@ +from __future__ import annotations + +from typing import Final + +from .....shared.tracing.steps import Engine, mapping +from ...fixtures import anthropic_response_body, anthropic_stream_events, json_response, sse_response +from ...models import GatewayRouteSpec, RouteFixture, TraceScenario, TraceSuite + +MAPPINGS: Final = ( + mapping(span="python_chat_gateway_route", python_frame=r"proxy_server\.py:\d+ chat_completion$"), + mapping(span="python_gateway_service", python_frame=r"ProxyBaseLLMRequestProcessing\.base_process_llm_request$"), + mapping(rust_span="chat_completions_gateway_route"), + mapping(rust_span="chat_completions"), + mapping(span="python_chat_entrypoint", python_frame=r"main\.py:\d+ a?completion$"), + mapping(span="python_provider_config", python_frame=r"ProviderConfigManager\.get_provider_chat_config$"), + mapping(rust_span="validate_environment", python_frame=r"(? RouteFixture: + return RouteFixture( + kwargs={ + "model_alias": "trace-model", + "provider_model": "anthropic/claude-sonnet-5", + "body": { + "model": "trace-model", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 16, + }, + }, + provider_responses=(json_response(anthropic_response_body()),), + ) + + +def _stream_fixture(engine: Engine, base_url: str) -> RouteFixture: + fixture: Final = _fixture(engine, base_url) + return fixture.with_body(stream=True).derive( + provider_responses=(sse_response(anthropic_stream_events()),), + ) + + +TRACE_SUITE: Final = TraceSuite( + route=GatewayRouteSpec("chat_completions"), + scenarios=( + TraceScenario(name="async-anthropic", fixture=_fixture, mappings=MAPPINGS, asynchronous=True), + TraceScenario( + name="async-anthropic-downstream-stream", + fixture=_stream_fixture, + mappings=(*MAPPINGS, *STREAM_MAPPINGS), + asynchronous=True, + ), + ), +) diff --git a/tests/rust-python-harness/strategies/trace_parity/gateway/execution.py b/tests/rust-python-harness/strategies/trace_parity/gateway/execution.py index 2bd3a50f39f..7a65edbef7f 100644 --- a/tests/rust-python-harness/strategies/trace_parity/gateway/execution.py +++ b/tests/rust-python-harness/strategies/trace_parity/gateway/execution.py @@ -12,8 +12,8 @@ from ....shared.parity.replay import replay_server from ....shared.tracing.native import TraceResponsePayload, native_trace_events from ....shared.tracing.profiler import FunctionTraceEvent, profile_python from ....shared.tracing.steps import Engine, PipelineProjection, pipeline_projection -from ..models import GatewayRouteSpec, RouteFixture, TraceExecutionFailure, TraceMode, TraceScenario -from ..reporting import TraceComparisonArtifact +from ..models import GatewayRouteSpec, RouteFixture, TraceEngine, TraceExecutionFailure, TraceScenario +from ..reporting import TraceArtifact class _GatewayResponsePayload(BaseModel): @@ -27,13 +27,20 @@ class _GatewayClient(Protocol): def post(self, url: str, *, json: object, headers: dict[str, str]) -> httpx.Response: ... -def _collect_python(fixture: RouteFixture) -> tuple[FunctionTraceEvent, ...]: - import litellm +_ROUTE_PATHS: Final = { + "messages": "/v1/messages", + "chat_completions": "/v1/chat/completions", + "responses": "/v1/responses", +} + + +def _collect_python(fixture: RouteFixture, route: GatewayRouteSpec) -> tuple[FunctionTraceEvent, ...]: from fastapi.testclient import TestClient - from litellm.proxy._types import UserAPIKeyAuth - from litellm.proxy.anthropic_endpoints.endpoints import user_api_key_auth + import litellm from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth provider_model: Final = cast(str, fixture.kwargs["provider_model"]) model_alias: Final = cast(str, fixture.kwargs["model_alias"]) @@ -60,7 +67,7 @@ def _collect_python(fixture: RouteFixture) -> tuple[FunctionTraceEvent, ...]: with profile_python(Path(litellm.__file__).parent, threads=True) as profiler: client: Final = cast(_GatewayClient, TestClient(proxy_server.app)) response: Final = client.post( - "/v1/messages", + _ROUTE_PATHS[route.route], json=fixture.kwargs["body"], headers={"authorization": "Bearer trace-key"}, ) @@ -75,15 +82,16 @@ def _collect_python(fixture: RouteFixture) -> tuple[FunctionTraceEvent, ...]: proxy_server.app.dependency_overrides[user_api_key_auth] = old_override -def _collect_rust(fixture: RouteFixture) -> tuple[FunctionTraceEvent, ...]: +def _collect_rust(fixture: RouteFixture, route: GatewayRouteSpec) -> tuple[FunctionTraceEvent, ...]: from litellm.rust_bridge import get_native_bridge bridge: Final[object | None] = get_native_bridge() trace: Final[object | None] = getattr(bridge, "_trace", None) if bridge is not None else None - gateway_messages: Final[object | None] = getattr(trace, "gateway_messages", None) - if gateway_messages is None or not callable(gateway_messages): - raise RuntimeError("native Rust trace bridge does not expose gateway_messages") - invoke_gateway: Final = cast(Callable[[str, str, str, object], Awaitable[object]], gateway_messages) + entrypoint: Final = f"gateway_{route.route}" + gateway_call: Final[object | None] = getattr(trace, entrypoint, None) + if gateway_call is None or not callable(gateway_call): + raise RuntimeError(f"native Rust trace bridge does not expose {entrypoint}") + invoke_gateway: Final = cast(Callable[[str, str, str, object], Awaitable[object]], gateway_call) async def invoke() -> object: return await invoke_gateway( @@ -101,7 +109,9 @@ def _collect_rust(fixture: RouteFixture) -> tuple[FunctionTraceEvent, ...]: return native_trace_events(payload) -def _collect(scenario: TraceScenario, engine: Engine) -> tuple[FunctionTraceEvent, ...] | TraceExecutionFailure: +def _collect( + route: GatewayRouteSpec, scenario: TraceScenario, engine: Engine +) -> tuple[FunctionTraceEvent, ...] | TraceExecutionFailure: try: with replay_server() as provider: base_fixture: Final = scenario.fixture(engine, provider.url) @@ -111,7 +121,7 @@ def _collect(scenario: TraceScenario, engine: Engine) -> tuple[FunctionTraceEven ) for response in fixture.provider_responses: provider.enqueue_response(response) - events: Final = _collect_python(fixture) if engine == "python" else _collect_rust(fixture) + events: Final = _collect_python(fixture, route) if engine == "python" else _collect_rust(fixture, route) provider.take_requests(len(fixture.provider_responses)) return events except Exception as error: @@ -122,9 +132,8 @@ def _projections( python_events: tuple[FunctionTraceEvent, ...], rust_events: tuple[FunctionTraceEvent, ...], scenario: TraceScenario, - mode: TraceMode, ) -> tuple[PipelineProjection, PipelineProjection, str | None]: - mappings: Final = scenario.mappings_for(mode) + mappings: Final = scenario.mappings try: return ( pipeline_projection("python", python_events, mappings), @@ -135,26 +144,26 @@ def _projections( return PipelineProjection(), PipelineProjection(), f"harness: {error}" -def execute_gateway_trace(route: GatewayRouteSpec, scenario: TraceScenario, mode: TraceMode) -> TraceComparisonArtifact: - mappings: Final = scenario.mappings_for(mode) - python_trace: Final = _collect(scenario, "python") - rust_trace: Final = _collect(scenario, "rust") +def execute_gateway_trace( + route: GatewayRouteSpec, + scenario: TraceScenario, + engine: TraceEngine = "both", +) -> TraceArtifact: + python_trace: Final = _collect(route, scenario, "python") if engine != "rust" else () + rust_trace: Final = _collect(route, scenario, "rust") if engine != "python" else () collection_python_error: Final = None if isinstance(python_trace, tuple) else f"python: {python_trace.message}" rust_error: Final = None if isinstance(rust_trace, tuple) else f"rust: {rust_trace.message}" python_events: Final = python_trace if isinstance(python_trace, tuple) else () rust_events: Final = rust_trace if isinstance(rust_trace, tuple) else () - python, rust, projection_error = _projections(python_events, rust_events, scenario, mode) + python, rust, projection_error = _projections(python_events, rust_events, scenario) python_error: Final = projection_error or collection_python_error - return TraceComparisonArtifact.from_traces( + return TraceArtifact.from_traces( + engine=engine, surface="gateway", sdk_function=route.route, scenario=scenario.name, - mode=mode, - mappings=mappings, - contract=scenario.contract, python=python.steps, rust=rust.steps, - python_unmatched=python.unmatched, python_error=python_error, rust_error=rust_error, ) diff --git a/tests/rust-python-harness/strategies/trace_parity/gateway/messages/case.py b/tests/rust-python-harness/strategies/trace_parity/gateway/messages/case.py index 30f51cee353..ca9c858f6b7 100644 --- a/tests/rust-python-harness/strategies/trace_parity/gateway/messages/case.py +++ b/tests/rust-python-harness/strategies/trace_parity/gateway/messages/case.py @@ -1,10 +1,9 @@ from __future__ import annotations -import json from typing import Final -from .....shared.parity.recorded_http import HttpHeader, RecordedHttpResponse from .....shared.tracing.steps import Engine, mapping +from ...fixtures import anthropic_response_body, anthropic_stream_events, json_response, sse_response from ...models import GatewayRouteSpec, RouteFixture, TraceScenario, TraceSuite @@ -49,24 +48,7 @@ def _fixture(_engine: Engine, provider: str) -> RouteFixture: "max_tokens": 16, }, }, - provider_responses=( - RecordedHttpResponse.from_bytes( - 200, - (HttpHeader(name="content-type", value="application/json"),), - json.dumps( - { - "id": "msg_trace", - "type": "message", - "role": "assistant", - "model": "claude-sonnet-5", - "content": [{"type": "text", "text": "hello"}], - "stop_reason": "end_turn", - "stop_sequence": None, - "usage": {"input_tokens": 2, "output_tokens": 3}, - } - ).encode(), - ), - ), + provider_responses=(json_response(anthropic_response_body()),), ) @@ -78,6 +60,13 @@ def _azure_fixture(engine: Engine, _base_url: str) -> RouteFixture: return _fixture(engine, "azure_ai") +def _stream_fixture(engine: Engine, _base_url: str) -> RouteFixture: + fixture: Final = _anthropic_fixture(engine, _base_url) + return fixture.with_body(stream=True).derive( + provider_responses=(sse_response(anthropic_stream_events()),), + ) + + ANTHROPIC_MAPPINGS: Final = ( *GATEWAY_MAPPINGS, mapping( @@ -100,7 +89,22 @@ AZURE_MAPPINGS: Final = ( TRACE_SUITE: Final = TraceSuite( route=GatewayRouteSpec("messages"), scenarios=( - TraceScenario(name="anthropic", fixture=_anthropic_fixture, mappings=ANTHROPIC_MAPPINGS, modes=("async",)), - TraceScenario(name="azure-ai", fixture=_azure_fixture, mappings=AZURE_MAPPINGS, modes=("async",)), + TraceScenario( + name="async-anthropic", fixture=_anthropic_fixture, mappings=ANTHROPIC_MAPPINGS, asynchronous=True + ), + TraceScenario(name="async-azure-ai", fixture=_azure_fixture, mappings=AZURE_MAPPINGS, asynchronous=True), + TraceScenario( + name="async-anthropic-downstream-stream", + fixture=_stream_fixture, + mappings=( + *ANTHROPIC_MAPPINGS, + mapping(span="python_upstream_stream", python_frame=r"AnthropicMessagesStreamingResponse\.__anext__$"), + mapping( + span="python_downstream_stream", python_frame=r"DataGenerator\.__anext__$|async_data_generator$" + ), + mapping(span="python_stream_callback", python_frame=r"Logging\.async_success_handler$"), + ), + asynchronous=True, + ), ), ) diff --git a/tests/rust-python-harness/strategies/trace_parity/gateway/responses/case.py b/tests/rust-python-harness/strategies/trace_parity/gateway/responses/case.py new file mode 100644 index 00000000000..52c1ebb391f --- /dev/null +++ b/tests/rust-python-harness/strategies/trace_parity/gateway/responses/case.py @@ -0,0 +1,63 @@ +from __future__ import annotations + +from typing import Final + +from .....shared.tracing.steps import Engine, mapping +from ...fixtures import json_response, responses_body, responses_stream_events, sse_response +from ...models import GatewayRouteSpec, RouteFixture, TraceScenario, TraceSuite + +MAPPINGS: Final = ( + mapping( + span="python_responses_gateway_route", python_frame=r"response_api_endpoints/endpoints\.py:\d+ responses_api$" + ), + mapping(span="python_gateway_service", python_frame=r"ProxyBaseLLMRequestProcessing\.base_process_llm_request$"), + mapping(rust_span="responses_gateway_route"), + mapping(span="python_responses", python_frame=r"responses/main\.py:\d+ a?responses$"), + mapping(span="python_provider_config", python_frame=r"ProviderConfigManager\.get_provider_responses_api_config$"), + mapping(rust_span="validate_environment", python_frame=r"OpenAIResponsesAPIConfig\.validate_environment$"), + mapping(rust_span="complete_url", python_frame=r"OpenAIResponsesAPIConfig\.get_complete_url$"), + mapping(rust_span="transform_request", python_frame=r"OpenAIResponsesAPIConfig\.transform_responses_api_request$"), + mapping(span="python_logging_pre_call", python_frame=r"Logging\.pre_call$"), + mapping(rust_span="http_request", python_frame=r"AsyncHTTPHandler\.post$|HTTPHandler\.post$"), + mapping(rust_span="transform_response", python_frame=r"OpenAIResponsesAPIConfig\.transform_response_api_response$"), + mapping(span="python_success_callback", python_frame=r"Logging\.async_success_handler$|Logging\.success_handler$"), +) + +STREAM_MAPPINGS: Final = ( + mapping(span="python_stream_iterator", python_frame=r"ResponsesAPIStreamingIterator\.__init__$"), + mapping(span="python_stream_next", python_frame=r"ResponsesAPIStreamingIterator\.__anext__$"), + mapping(span="python_stream_transform", python_frame=r"OpenAIResponsesAPIConfig\.transform_streaming_response$"), + mapping(span="python_downstream_stream", python_frame=r"DataGenerator\.__anext__$|async_data_generator$"), +) + + +def _fixture(_engine: Engine, _base_url: str) -> RouteFixture: + return RouteFixture( + kwargs={ + "model_alias": "trace-model", + "provider_model": "openai/gpt-5", + "body": {"model": "trace-model", "input": "hello"}, + }, + provider_responses=(json_response(responses_body()),), + ) + + +def _stream_fixture(engine: Engine, base_url: str) -> RouteFixture: + fixture: Final = _fixture(engine, base_url) + return fixture.with_body(stream=True).derive( + provider_responses=(sse_response(responses_stream_events()),), + ) + + +TRACE_SUITE: Final = TraceSuite( + route=GatewayRouteSpec("responses"), + scenarios=( + TraceScenario(name="async-openai", fixture=_fixture, mappings=MAPPINGS, asynchronous=True), + TraceScenario( + name="async-openai-downstream-stream", + fixture=_stream_fixture, + mappings=(*MAPPINGS, *STREAM_MAPPINGS), + asynchronous=True, + ), + ), +) diff --git a/tests/rust-python-harness/strategies/trace_parity/models.py b/tests/rust-python-harness/strategies/trace_parity/models.py index 04659b25382..078af2b316d 100644 --- a/tests/rust-python-harness/strategies/trace_parity/models.py +++ b/tests/rust-python-harness/strategies/trace_parity/models.py @@ -1,22 +1,45 @@ from __future__ import annotations -from collections.abc import Callable +from collections.abc import Callable, Mapping from dataclasses import dataclass -from typing import Final, Literal, TypeAlias +from typing import Final, Literal, TypeAlias, cast -from ...shared.parity.recorded_http import RecordedHttpResponse +from ...shared.parity.recorded_http import RecordedResponse from ...shared.reporting.models import SdkFunction -from ...shared.tracing.steps import Engine, TraceContract, TraceMapping +from ...shared.tracing.steps import Engine, TraceMapping -TraceMode = Literal["sync", "async"] +TraceEngine = Literal["python", "rust", "both"] TraceFailureSource = Literal["python", "rust", "harness"] @dataclass(frozen=True, slots=True) class RouteFixture: kwargs: dict[str, object] - provider_responses: tuple[RecordedHttpResponse, ...] + provider_responses: tuple[RecordedResponse, ...] expected_failure: bool = False + consume_stream: bool = False + + def derive( + self, + *, + kwargs: Mapping[str, object] | None = None, + provider_responses: tuple[RecordedResponse, ...] | None = None, + expected_failure: bool | None = None, + consume_stream: bool | None = None, + ) -> RouteFixture: + return RouteFixture( + kwargs={**self.kwargs, **(kwargs or {})}, + provider_responses=self.provider_responses if provider_responses is None else provider_responses, + expected_failure=self.expected_failure if expected_failure is None else expected_failure, + consume_stream=self.consume_stream if consume_stream is None else consume_stream, + ) + + def with_body(self, **updates: object) -> RouteFixture: + raw_body: Final = self.kwargs.get("body") + if not isinstance(raw_body, dict): + raise ValueError("route fixture does not contain an object body") + body: Final = cast(dict[str, object], raw_body) + return self.derive(kwargs={"body": {**body, **updates}}) @dataclass(frozen=True, slots=True) @@ -40,14 +63,7 @@ class TraceScenario: name: str fixture: Callable[[Engine, str], RouteFixture] mappings: tuple[TraceMapping, ...] - modes: tuple[TraceMode, ...] = ("sync", "async") - contract: TraceContract = TraceContract() - sync_mappings: tuple[TraceMapping, ...] | None = None - async_mappings: tuple[TraceMapping, ...] | None = None - - def mappings_for(self, mode: TraceMode) -> tuple[TraceMapping, ...]: - selected: Final = self.async_mappings if mode == "async" else self.sync_mappings - return self.mappings if selected is None else selected + asynchronous: bool @dataclass(frozen=True, slots=True) diff --git a/tests/rust-python-harness/strategies/trace_parity/reporting.py b/tests/rust-python-harness/strategies/trace_parity/reporting.py index 9c5bf9e88cd..e7c07ef9c0f 100644 --- a/tests/rust-python-harness/strategies/trace_parity/reporting.py +++ b/tests/rust-python-harness/strategies/trace_parity/reporting.py @@ -1,31 +1,24 @@ from __future__ import annotations import os -import re import sys from collections.abc import Sequence -from typing import Final, Literal +from typing import Final from pydantic import BaseModel, ConfigDict, ValidationError from ...shared.reporting.models import SURFACES, CaseResult, RunStatus, SdkFunction, Surface from ...shared.reporting.rendering import ReportSection from ...shared.reporting.strategy import NotImplementedCaseSpec, SkippedCaseSpec -from ...shared.tracing.steps import ( - PipelineStep, - TraceContract, - TraceDiff, - TraceMapping, - trace_depths, - trace_diff, -) +from ...shared.tracing.steps import PipelineStep, trace_depths +from .models import TraceEngine -TRACE_COMPARISON_ARTIFACT: Final = "trace_comparison" +TRACE_ARTIFACT: Final = "trace" TRACE_PARITY_HINT: Final = ( "rebuild the native bridge with the trace-parity feature, e.g. `uvx maturin develop --features trace-parity`" ) -_COLORS: Final[dict[str, str]] = {"green": "32", "yellow": "33", "red": "31", "cyan": "36"} +_COLORS: Final[dict[str, str]] = {"yellow": "33", "red": "31", "cyan": "36"} _RESET: Final = "\033[0m" @@ -47,26 +40,15 @@ class TraceEventArtifact(BaseModel): return PipelineStep(self.id, self.parent_id, self.span, self.raw) -class TraceMappingArtifact(BaseModel): - model_config = ConfigDict(frozen=True, extra="forbid") - - span: str - python: str | None - rust: str | None - - -class TraceComparisonArtifact(BaseModel): +class TraceArtifact(BaseModel): model_config = ConfigDict(frozen=True, extra="forbid") + engine: TraceEngine = "both" surface: Surface sdk_function: SdkFunction scenario: str - mode: Literal["sync", "async"] - mappings: tuple[TraceMappingArtifact, ...] python: tuple[TraceEventArtifact, ...] rust: tuple[TraceEventArtifact, ...] - python_unmatched: int - unordered_children_of: frozenset[str] python_error: str | None = None rust_error: str | None = None @@ -74,41 +56,27 @@ class TraceComparisonArtifact(BaseModel): def from_traces( cls, *, + engine: TraceEngine = "both", surface: Surface, sdk_function: SdkFunction, scenario: str, - mode: Literal["sync", "async"], - mappings: Sequence[TraceMapping], - contract: TraceContract, python: Sequence[PipelineStep], rust: Sequence[PipelineStep], - python_unmatched: int, python_error: str | None = None, rust_error: str | None = None, - ) -> TraceComparisonArtifact: + ) -> TraceArtifact: return cls( + engine=engine, surface=surface, sdk_function=sdk_function, scenario=scenario, - mode=mode, - mappings=tuple( - TraceMappingArtifact( - span=item.span, - python=item.python.pattern if item.python else None, - rust=item.rust, - ) - for item in mappings - ), python=tuple( TraceEventArtifact(id=step.id, parent_id=step.parent_id, span=step.span, raw=step.raw) for step in python ), rust=tuple( - TraceEventArtifact(id=step.id, parent_id=step.parent_id, span=step.span, raw=step.raw) - for step in rust + TraceEventArtifact(id=step.id, parent_id=step.parent_id, span=step.span, raw=step.raw) for step in rust ), - python_unmatched=python_unmatched, - unordered_children_of=contract.unordered_children_of, python_error=python_error, rust_error=rust_error, ) @@ -119,32 +87,9 @@ class TraceComparisonArtifact(BaseModel): def rust_steps(self) -> tuple[PipelineStep, ...]: return tuple(event.step() for event in self.rust) - def diff(self) -> TraceDiff: - return trace_diff( - self.python_steps(), - self.rust_steps(), - tuple( - TraceMapping( - item.span, - re.compile(item.python) if item.python is not None else None, - item.rust, - ) - for item in self.mappings - ), - TraceContract(self.unordered_children_of), - ) - - def exact_match(self) -> bool: - return self.diff().matches - def has_errors(self) -> bool: return self.python_error is not None or self.rust_error is not None - def contract_matches(self) -> bool: - if self.has_errors(): - return False - return self.diff().matches - def _split_raw(raw: str) -> tuple[str, str]: location, separator, name = raw.partition(" ") @@ -153,72 +98,28 @@ def _split_raw(raw: str) -> tuple[str, str]: return raw, "" -def _python_line(index: int, step: PipelineStep, depth: int, exclusive: frozenset[str]) -> str: +def _python_line(index: int, step: PipelineStep, depth: int) -> str: name: Final = _split_raw(step.raw)[0] location: Final = _split_raw(step.raw)[1] suffix: Final = f" ({location})" if location else "" - marker: Final = " [python only]" if step.span in exclusive else "" - return _paint(f"{index} {' ' * depth}{name}{suffix}{marker}", "cyan") + return _paint(f"{index} {' ' * depth}{name}{suffix}", "cyan") -def _python_lines(steps: tuple[PipelineStep, ...], exclusive: frozenset[str]) -> str: +def _python_lines(steps: tuple[PipelineStep, ...]) -> str: depths: Final = trace_depths(steps) - lines: Final = tuple( - _python_line(index, step, depths[step.id], exclusive) for index, step in enumerate(steps, start=1) - ) + lines: Final = tuple(_python_line(index, step, depths[step.id]) for index, step in enumerate(steps, start=1)) return f"{_paint('PYTHON', 'cyan')} ({len(steps)} steps)\n" + ("\n".join(lines) if lines else "(empty)") -def _python_references(steps: tuple[PipelineStep, ...]) -> dict[tuple[str, int], str]: - references: dict[tuple[str, int], str] = {} - occurrences: dict[str, int] = {} - for index, step in enumerate(steps, start=1): - name = _split_raw(step.raw)[0] - occurrence = occurrences.get(step.span, 0) + 1 - occurrences[step.span] = occurrence - references[(step.span, occurrence)] = f"{index} {name}" - return references - - -def _rust_line( - step: PipelineStep, - depth: int, - occurrence: int, - references: dict[tuple[str, int], str], -) -> str: - span: Final = _paint(step.span, "yellow") - key: Final = (step.span, occurrence) - reference: Final = ( - _paint(references[key], "cyan") if key in references else _paint("[rust only]", "yellow") - ) - suffix: Final = f"#{occurrence}" if occurrence > 1 else "" - return f"{' ' * depth}{span}{suffix} -> {reference}" - - -def _rust_lines(steps: tuple[PipelineStep, ...], references: dict[tuple[str, int], str]) -> str: +def _rust_lines(steps: tuple[PipelineStep, ...]) -> str: depths: Final = trace_depths(steps) - occurrences: dict[str, int] = {} - lines: list[str] = [] - for step in steps: - occurrence = occurrences.get(step.span, 0) + 1 - occurrences[step.span] = occurrence - lines.append(_rust_line(step, depths[step.id], occurrence, references)) + lines: Final = tuple( + _paint(f"{index} {' ' * depths[step.id]}{step.span}", "yellow") for index, step in enumerate(steps, 1) + ) return f"{_paint('RUST', 'yellow')} ({len(steps)} steps)\n" + ("\n".join(lines) if lines else "(empty)") -def _state_text(state: str, *, good: bool) -> str: - return _paint(state, "green" if good else "red") - - -def _contract_line(artifact: TraceComparisonArtifact) -> str: - matches: Final = artifact.contract_matches() - status: Final = _state_text("PASS" if matches else "FAIL", good=matches) - if artifact.python_error or artifact.rust_error: - return f"Contract: {status}" - return f"Contract: {status}" - - -def _error_lines(artifact: TraceComparisonArtifact) -> tuple[str, ...]: +def _error_lines(artifact: TraceArtifact) -> tuple[str, ...]: lines: list[str] = [] for engine, error in (("Python", artifact.python_error), ("Rust", artifact.rust_error)): if error is None: @@ -229,68 +130,20 @@ def _error_lines(artifact: TraceComparisonArtifact) -> tuple[str, ...]: return tuple(lines) -def _unseen_mappings( - artifact: TraceComparisonArtifact, - python: tuple[PipelineStep, ...], - rust: tuple[PipelineStep, ...], -) -> tuple[str, ...]: - return artifact.diff().missing_mappings - - -def _comparison_status_lines( - artifact: TraceComparisonArtifact, - python: tuple[PipelineStep, ...], - rust: tuple[PipelineStep, ...], -) -> tuple[str, ...]: - diff: Final = artifact.diff() - exact_match: Final = artifact.exact_match() - if artifact.has_errors(): - return (*_error_lines(artifact), _contract_line(artifact)) - unseen: Final = _unseen_mappings(artifact, python, rust) - unseen_line: Final[tuple[str, ...]] = (f"Unseen mappings: {', '.join(unseen)}",) if unseen else () - drift_lines: Final[tuple[str, ...]] = ( - (_state_text("Same steps, order, and nesting", good=True),) - if exact_match - else ( - _paint(f"Python only: {', '.join(diff.python_only) or 'none'}", "cyan"), - _paint(f"Rust only: {', '.join(diff.rust_only) or 'none'}", "yellow"), - f"First difference: {diff.first_difference or 'none'}", - f"Python frames outside mapping: {artifact.python_unmatched}", - ) - ) - return ( - f"Trace: {_state_text('MATCH' if exact_match else 'DRIFT', good=exact_match)}", - *drift_lines, - *unseen_line, - _contract_line(artifact), - ) - - -def _render_comparison(artifact: TraceComparisonArtifact) -> str: - python: Final = artifact.python_steps() - rust: Final = artifact.rust_steps() - diff: Final = artifact.diff() - python_exclusive: Final = frozenset(item.span for item in artifact.mappings if item.rust is None) - status_lines: Final = _comparison_status_lines(artifact, python, rust) - return "\n\n".join( - ( - _python_lines(python, python_exclusive | frozenset(diff.python_only)), - _rust_lines(rust, _python_references(python)), - "\n".join(status_lines), - ) - ) - - -def _mode(nodeid: str) -> str: - if "[" in nodeid: - return nodeid.rsplit("[", 1)[-1].removesuffix("]") - head, _, tail = nodeid.rpartition(":") - return tail if head else "unknown mode" +def _render_trace(artifact: TraceArtifact) -> str: + traces: tuple[str, ...] + if artifact.engine == "python": + traces = (_python_lines(artifact.python_steps()),) + elif artifact.engine == "rust": + traces = (_rust_lines(artifact.rust_steps()),) + else: + traces = (_python_lines(artifact.python_steps()), _rust_lines(artifact.rust_steps())) + return "\n\n".join((*traces, *_error_lines(artifact))) def _scenario(nodeid: str) -> str: parts: Final = nodeid.split(":") - return parts[-2] if len(parts) >= 5 else "default" + return parts[-1] if len(parts) >= 4 else "default" def _unavailable(status: RunStatus) -> str: @@ -299,20 +152,20 @@ def _unavailable(status: RunStatus) -> str: def _render_artifact(body: str) -> str: try: - artifact: Final = TraceComparisonArtifact.model_validate_json(body) + artifact: Final = TraceArtifact.model_validate_json(body) except ValidationError as error: - return f"Trace comparison artifact is invalid: {error}" - return _render_comparison(artifact) + return f"Trace artifact is invalid: {error}" + return _render_trace(artifact) -def _mode_section(result: CaseResult, nodeid: str, status: RunStatus) -> str: +def _scenario_section(result: CaseResult, nodeid: str, status: RunStatus) -> str: artifacts: Final = tuple( - artifact for artifact in result.artifacts.get(nodeid, ()) if artifact.kind == TRACE_COMPARISON_ARTIFACT + artifact for artifact in result.artifacts.get(nodeid, ()) if artifact.kind == TRACE_ARTIFACT ) body: Final = ( "\n\n".join(_render_artifact(artifact.body) for artifact in artifacts) if artifacts else _unavailable(status) ) - label: Final = f"Scenario: {_scenario(nodeid)} / Mode: {_mode(nodeid)}" + label: Final = f"Scenario: {_scenario(nodeid)}" return f"{label}\n{'-' * len(label)}\n\n{body}" @@ -321,7 +174,7 @@ def _case_block(result: CaseResult) -> str: outcomes: Final = tuple(result.outcomes.items()) or ( (nodeid, RunStatus.NOT_RUN) for nodeid in sorted(result.collected) ) - sections: Final = tuple(_mode_section(result, nodeid, status) for nodeid, status in outcomes) + sections: Final = tuple(_scenario_section(result, nodeid, status) for nodeid, status in outcomes) return "\n\n".join((f"{header}\n{'=' * len(header)}", *sections)) @@ -357,11 +210,11 @@ def _surface_section(surface: Surface, results: Sequence[CaseResult]) -> ReportS *((not_implemented,) if not_implemented else ()), *((skipped,) if skipped else ()), ) - return ReportSection(f"{surface.upper()} trace comparisons", blocks or ("No runnable trace comparisons",)) + return ReportSection(f"{surface.upper()} traces", blocks or ("No runnable traces",)) def render_trace_results(results: Sequence[CaseResult]) -> tuple[ReportSection, ...]: sections: Final = tuple( section for surface in SURFACES if (section := _surface_section(surface, results)) is not None ) - return sections or (ReportSection("Trace comparisons", ("No trace comparisons selected",)),) + return sections or (ReportSection("Traces", ("No traces selected",)),) diff --git a/tests/rust-python-harness/strategies/trace_parity/runner.py b/tests/rust-python-harness/strategies/trace_parity/runner.py index b78a3c7da3f..706e054bf52 100644 --- a/tests/rust-python-harness/strategies/trace_parity/runner.py +++ b/tests/rust-python-harness/strategies/trace_parity/runner.py @@ -4,13 +4,20 @@ import importlib from collections.abc import Sequence from pathlib import Path from time import monotonic -from typing import Final +from typing import Final, cast +from ...shared.native_build import ensure_trace_bridge from ...shared.reporting.models import CaseResult, HarnessCase, HarnessRun, ResultArtifact, RunStatus, Surface from ...shared.reporting.strategy import ModuleCaseSpec, UpdateCallback -from ...shared.native_build import ensure_trace_bridge -from .models import GatewayRouteSpec, RouteSpec, TraceExecutionFailure, TraceMode, TraceScenario, TraceSuite -from .reporting import TRACE_COMPARISON_ARTIFACT, TraceComparisonArtifact +from .models import ( + GatewayRouteSpec, + RouteSpec, + TraceEngine, + TraceExecutionFailure, + TraceScenario, + TraceSuite, +) +from .reporting import TRACE_ARTIFACT, TraceArtifact from .sdk.execution import execute_trace @@ -32,15 +39,13 @@ def validate_trace_suite(suite: TraceSuite, harness_case: HarnessCase) -> str | names: Final = tuple(scenario.name for scenario in suite.scenarios) if not names or len(names) != len(set(names)) or any(not name or ":" in name for name in names): return "scenario names must be non-empty, unique, and colon-free" - invalid_modes: Final = tuple( + invalid_names: Final = tuple( scenario.name for scenario in suite.scenarios - if not scenario.modes - or len(scenario.modes) != len(set(scenario.modes)) - or any(mode not in {"sync", "async"} for mode in scenario.modes) + if not scenario.name.startswith("async-" if scenario.asynchronous else "sync-") ) - if invalid_modes: - return f"scenarios must use non-empty, unique sync/async modes: {', '.join(invalid_modes)}" + if invalid_names: + return f"scenario names must start with sync- or async-: {', '.join(invalid_names)}" surface: Final = harness_case.surface if surface == "sdk" and not isinstance(suite.route, RouteSpec): return "must use RouteSpec for the sdk surface" @@ -57,15 +62,14 @@ def scenario_nodeids( trace_suite: TraceSuite, harness_case: HarnessCase, selected_scenarios: frozenset[str] = frozenset(), -) -> tuple[tuple[TraceScenario, TraceMode, str], ...]: +) -> tuple[tuple[TraceScenario, str], ...]: surface: Final = harness_case.surface if surface is None: return () return tuple( - (scenario, mode, f"trace:{surface}:{harness_case.sdk_function}:{scenario.name}:{mode}") + (scenario, f"trace:{surface}:{harness_case.sdk_function}:{scenario.name}") for scenario in trace_suite.scenarios if not selected_scenarios or scenario.name in selected_scenarios - for mode in scenario.modes ) @@ -77,54 +81,49 @@ def _record_setup_failure(run: HarnessRun, case: HarnessCase, message: str, stag run.failures.append((nodeid, message)) -def run_trace_mode( +def run_trace_scenario( run: HarnessRun, result: CaseResult, trace_suite: TraceSuite, scenario: TraceScenario, - mode: TraceMode, surface: Surface, nodeid: str, on_update: UpdateCallback, + engine: TraceEngine = "both", ) -> None: started_at: Final = monotonic() - comparison: Final = _execute_mode(trace_suite, scenario, mode, surface) + trace: Final = _execute_scenario(trace_suite, scenario, surface, engine) duration: Final = monotonic() - started_at - if isinstance(comparison, TraceExecutionFailure): + if isinstance(trace, TraceExecutionFailure): result.record(nodeid, RunStatus.ERROR, duration) - run.failures.append((nodeid, comparison.message)) + run.failures.append((nodeid, trace.message)) on_update(run) return - artifact: Final = ResultArtifact(TRACE_COMPARISON_ARTIFACT, comparison.model_dump_json()) - if comparison.has_errors(): + artifact: Final = ResultArtifact(TRACE_ARTIFACT, trace.model_dump_json()) + if trace.has_errors(): result.record(nodeid, RunStatus.ERROR, duration, (artifact,)) - run.failures.append( - (nodeid, "\n".join(error for error in (comparison.python_error, comparison.rust_error) if error)) - ) + run.failures.append((nodeid, "\n".join(error for error in (trace.python_error, trace.rust_error) if error))) else: - status: Final = RunStatus.PASSED if comparison.contract_matches() else RunStatus.FAILED - result.record(nodeid, status, duration, (artifact,)) - if status is RunStatus.FAILED: - run.failures.append((nodeid, "trace contract mismatch; see the rendered comparison")) + result.record(nodeid, RunStatus.PASSED, duration, (artifact,)) on_update(run) -def _execute_mode( +def _execute_scenario( trace_suite: TraceSuite, scenario: TraceScenario, - mode: TraceMode, surface: Surface, -) -> TraceComparisonArtifact | TraceExecutionFailure: + engine: TraceEngine, +) -> TraceArtifact | TraceExecutionFailure: route: Final = trace_suite.route if isinstance(route, GatewayRouteSpec): if surface != "gateway": return TraceExecutionFailure("harness", "gateway route cannot run on the sdk surface") from .gateway.execution import execute_gateway_trace - return execute_gateway_trace(route, scenario, mode) + return execute_gateway_trace(route, scenario, engine) if surface != "sdk": return TraceExecutionFailure("harness", "sdk route cannot run on the gateway surface") - return execute_trace(route, scenario, mode, surface) + return execute_trace(route, scenario, surface, engine) def _run_case( @@ -132,6 +131,7 @@ def _run_case( harness_case: HarnessCase, selected_scenarios: frozenset[str], on_update: UpdateCallback, + engine: TraceEngine, ) -> None: result: Final = run.results[harness_case.key] spec: Final = harness_case.spec @@ -146,15 +146,29 @@ def _run_case( on_update(run) return nodeids: Final = scenario_nodeids(trace_suite, harness_case, selected_scenarios) - result.collected.update(nodeid for _, _, nodeid in nodeids) + result.collected.update(nodeid for _, nodeid in nodeids) if not nodeids: result.status = RunStatus.SKIPPED on_update(run) return result.status = RunStatus.RUNNING on_update(run) - for scenario, mode, nodeid in nodeids: - run_trace_mode(run, result, trace_suite, scenario, mode, surface, nodeid, on_update) + for scenario, nodeid in nodeids: + run_trace_scenario(run, result, trace_suite, scenario, surface, nodeid, on_update, engine) + + +def runner_selection(runner_args: Sequence[str]) -> tuple[frozenset[str], TraceEngine]: + engine: TraceEngine = "both" + scenarios: list[str] = [] + for argument in runner_args: + if argument.startswith("--engine="): + value = argument.removeprefix("--engine=") + if value not in {"python", "rust"}: + raise ValueError(f"invalid trace engine: {value}") + engine = cast(TraceEngine, value) + else: + scenarios.append(argument) + return frozenset(scenarios), engine def run_trace_cases( @@ -163,10 +177,10 @@ def run_trace_cases( on_update: UpdateCallback, runner_args: Sequence[str] = (), ) -> tuple[int, HarnessRun]: - selected_scenarios: Final = frozenset(runner_args) + selected_scenarios, engine = runner_selection(runner_args) run: Final = HarnessRun.from_cases(cases) runnable_cases: Final = tuple(case for case in cases if isinstance(case.spec, ModuleCaseSpec)) - bridge_error: Final = ensure_trace_bridge(repo_root) if runnable_cases else None + bridge_error: Final = ensure_trace_bridge(repo_root) if runnable_cases and engine != "python" else None if bridge_error is not None: for harness_case in runnable_cases: _record_setup_failure(run, harness_case, bridge_error, "bridge") @@ -174,7 +188,7 @@ def run_trace_cases( on_update(run) return 1, run for harness_case in cases: - _run_case(run, harness_case, selected_scenarios, on_update) + _run_case(run, harness_case, selected_scenarios, on_update, engine) run.finished_at = monotonic() on_update(run) failed: Final = any( diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/chat_completions/case.py b/tests/rust-python-harness/strategies/trace_parity/sdk/chat_completions/case.py index 6be5afd60d6..1221f237570 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/chat_completions/case.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/chat_completions/case.py @@ -1,10 +1,15 @@ from __future__ import annotations -import json from typing import Final -from .....shared.parity.recorded_http import HttpHeader, RecordedHttpResponse from .....shared.tracing.steps import Engine, mapping +from ...fixtures import ( + anthropic_response_body, + anthropic_stream_events, + aws_event_stream_response, + json_response, + sse_response, +) from ...models import RouteFixture, RouteSpec, TraceScenario, TraceSuite COMMON_MAPPINGS: Final = ( @@ -24,6 +29,22 @@ COMMON_MAPPINGS: Final = ( mapping(rust_span="execute_chat_completions_provider_call"), mapping(rust_span="http_request", python_frame=r"AsyncHTTPHandler\.post$|HTTPHandler\.post$"), mapping(rust_span="transform_response", python_frame=r"(? RouteFixture: - response: Final = json.dumps( - { - "id": "msg_trace", - "type": "message", - "role": "assistant", - "model": "claude-sonnet-5", - "content": [{"type": "text", "text": "hello"}], - "stop_reason": "end_turn", - "stop_sequence": None, - "usage": {"input_tokens": 2, "output_tokens": 3}, - } - ).encode() return RouteFixture( kwargs={ "model": "anthropic/claude-sonnet-5", "messages": [{"role": "user", "content": "hello"}], **({"optional_params": {"max_tokens": 16}} if engine == "rust" else {"max_tokens": 16}), }, - provider_responses=( - RecordedHttpResponse.from_bytes( - 200, (HttpHeader(name="content-type", value="application/json"),), response - ), - ), + provider_responses=(json_response(anthropic_response_body()),), ) def _bedrock_fixture(engine: Engine, _base_url: str) -> RouteFixture: - response: Final = json.dumps( - { - "output": {"message": {"role": "assistant", "content": [{"text": "hello"}]}}, - "stopReason": "end_turn", - "usage": {"inputTokens": 2, "outputTokens": 3, "totalTokens": 5}, - "metrics": {"latencyMs": 1}, - } - ).encode() + response: Final[dict[str, object]] = { + "output": {"message": {"role": "assistant", "content": [{"text": "hello"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 2, "outputTokens": 3, "totalTokens": 5}, + "metrics": {"latencyMs": 1}, + } credentials: Final = { "aws_access_key_id": "test-access", "aws_secret_access_key": "test-secret", @@ -94,11 +97,60 @@ def _bedrock_fixture(engine: Engine, _base_url: str) -> RouteFixture: else {**credentials, "max_tokens": 16} ), }, + provider_responses=(json_response(response),), + ) + + +def _anthropic_stream_fixture(engine: Engine, _base_url: str) -> RouteFixture: + fixture: Final = _anthropic_fixture(engine, _base_url) + return fixture.derive( + kwargs={"stream": True}, + provider_responses=(sse_response(anthropic_stream_events()),), + consume_stream=True, + ) + + +def _bedrock_stream_fixture(engine: Engine, _base_url: str) -> RouteFixture: + fixture: Final = _bedrock_fixture(engine, _base_url) + events: Final[tuple[dict[str, object], ...]] = ( + {"messageStart": {"role": "assistant"}}, + {"contentBlockStart": {"contentBlockIndex": 0, "start": {}}}, + {"contentBlockDelta": {"contentBlockIndex": 0, "delta": {"text": "hello"}}}, + {"contentBlockStop": {"contentBlockIndex": 0}}, + {"messageStop": {"stopReason": "end_turn"}}, + {"metadata": {"usage": {"inputTokens": 2, "outputTokens": 1, "totalTokens": 3}}}, + ) + return fixture.derive( + kwargs={"stream": True}, + provider_responses=(aws_event_stream_response(events),), + consume_stream=True, + ) + + +def _provider_error_fixture(engine: Engine, _base_url: str) -> RouteFixture: + fixture: Final = _anthropic_fixture(engine, _base_url) + return fixture.derive( provider_responses=( - RecordedHttpResponse.from_bytes( - 200, (HttpHeader(name="content-type", value="application/json"),), response + json_response( + {"type": "error", "error": {"type": "invalid_request_error", "message": "bad request"}}, + status=400, ), ), + expected_failure=True, + ) + + +def _stream_error_fixture(engine: Engine, base_url: str) -> RouteFixture: + fixture: Final = _anthropic_fixture(engine, base_url) + events: Final = ( + anthropic_stream_events()[0], + ("error", {"type": "error", "error": {"type": "overloaded_error", "message": "overloaded"}}), + ) + return fixture.derive( + kwargs={"stream": True}, + provider_responses=(sse_response(events),), + expected_failure=True, + consume_stream=True, ) @@ -132,18 +184,64 @@ TRACE_SUITE: Final = TraceSuite( route=SPEC, scenarios=( TraceScenario( - name="anthropic", + name="sync-anthropic", fixture=_anthropic_fixture, - mappings=COMMON_MAPPINGS, - sync_mappings=SYNC_MAPPINGS, - async_mappings=ASYNC_MAPPINGS, + mappings=SYNC_MAPPINGS, + asynchronous=False, ), TraceScenario( - name="bedrock", + name="async-anthropic", + fixture=_anthropic_fixture, + mappings=ASYNC_MAPPINGS, + asynchronous=True, + ), + TraceScenario( + name="sync-anthropic-stream", + fixture=_anthropic_stream_fixture, + mappings=(*SYNC_MAPPINGS, *STREAM_MAPPINGS), + asynchronous=False, + ), + TraceScenario( + name="async-anthropic-stream", + fixture=_anthropic_stream_fixture, + mappings=(*ASYNC_MAPPINGS, *STREAM_MAPPINGS), + asynchronous=True, + ), + TraceScenario( + name="async-anthropic-provider-error", + fixture=_provider_error_fixture, + mappings=(*ASYNC_MAPPINGS, *FAILURE_MAPPINGS), + asynchronous=True, + ), + TraceScenario( + name="async-anthropic-stream-error", + fixture=_stream_error_fixture, + mappings=(*ASYNC_MAPPINGS, *STREAM_MAPPINGS, *FAILURE_MAPPINGS), + asynchronous=True, + ), + TraceScenario( + name="sync-bedrock", fixture=_bedrock_fixture, - mappings=BEDROCK_COMMON_MAPPINGS, - sync_mappings=BEDROCK_SYNC_MAPPINGS, - async_mappings=BEDROCK_ASYNC_MAPPINGS, + mappings=BEDROCK_SYNC_MAPPINGS, + asynchronous=False, + ), + TraceScenario( + name="async-bedrock", + fixture=_bedrock_fixture, + mappings=BEDROCK_ASYNC_MAPPINGS, + asynchronous=True, + ), + TraceScenario( + name="sync-bedrock-event-stream", + fixture=_bedrock_stream_fixture, + mappings=(*BEDROCK_SYNC_MAPPINGS, *STREAM_MAPPINGS), + asynchronous=False, + ), + TraceScenario( + name="async-bedrock-event-stream", + fixture=_bedrock_stream_fixture, + mappings=(*BEDROCK_ASYNC_MAPPINGS, *STREAM_MAPPINGS), + asynchronous=True, ), ), ) diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py b/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py index eb6c9233565..c249d09fa73 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py @@ -1,7 +1,7 @@ from __future__ import annotations import asyncio -from collections.abc import Awaitable +from collections.abc import AsyncIterable, Awaitable, Iterable from dataclasses import dataclass from pathlib import Path from typing import Final, Protocol, cast @@ -11,8 +11,8 @@ from ....shared.reporting.models import Surface from ....shared.tracing.native import TraceResponsePayload, native_trace_events from ....shared.tracing.profiler import FunctionTraceEvent, profile_python from ....shared.tracing.steps import Engine, pipeline_projection -from ..models import RouteFixture, RouteSpec, TraceExecutionFailure, TraceMode, TraceScenario -from ..reporting import TraceComparisonArtifact +from ..models import RouteFixture, RouteSpec, TraceEngine, TraceExecutionFailure, TraceScenario +from ..reporting import TraceArtifact class SdkCall(Protocol): @@ -25,10 +25,20 @@ class _CollectedTrace: error: str | None = None -def _invoke(function: SdkCall, kwargs: dict[str, object], *, asynchronous: bool) -> object: +def _invoke( + function: SdkCall, + kwargs: dict[str, object], + *, + asynchronous: bool, + consume_stream: bool = False, +) -> object: async def invoke_async() -> object: try: - return await cast(Awaitable[object], function(**kwargs)) + response: Final = await cast(Awaitable[object], function(**kwargs)) + if consume_stream and isinstance(response, AsyncIterable): + stream = cast(AsyncIterable[object], response) + return tuple([item async for item in stream]) + return response finally: await asyncio.sleep(0) from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER @@ -38,7 +48,10 @@ def _invoke(function: SdkCall, kwargs: dict[str, object], *, asynchronous: bool) if asynchronous: return asyncio.run(invoke_async()) - return function(**kwargs) + response: Final = function(**kwargs) + if consume_stream and isinstance(response, Iterable): + return tuple(cast(Iterable[object], response)) + return response def _entrypoint(spec: RouteSpec, engine: Engine, *, asynchronous: bool) -> SdkCall | TraceExecutionFailure: @@ -62,14 +75,6 @@ def _entrypoint(spec: RouteSpec, engine: Engine, *, asynchronous: bool) -> SdkCa return cast(SdkCall, getattr(owner, spec.python_entrypoints[int(asynchronous)])) -def _python_invocation_error(function: SdkCall, kwargs: dict[str, object], *, asynchronous: bool) -> str | None: - try: - _invoke(function, kwargs, asynchronous=asynchronous) - except Exception as error: - return f"{type(error).__name__}: {error}" - return None - - def _collect( function: SdkCall, fixture: RouteFixture, @@ -83,8 +88,19 @@ def _collect( return _CollectedTrace(native_trace_events(payload), payload.error) import litellm - with profile_python(Path(litellm.__file__).parent, threads=True) as profiler: - error: Final = _python_invocation_error(function, kwargs, asynchronous=asynchronous) + previous_suppress_debug_info: Final = litellm.suppress_debug_info + try: + if fixture.expected_failure: + litellm.suppress_debug_info = True + with profile_python(Path(litellm.__file__).parent, threads=True) as profiler: + error: str | None + try: + _invoke(function, kwargs, asynchronous=asynchronous, consume_stream=fixture.consume_stream) + error = None + except Exception as caught: + error = f"{type(caught).__name__}: {caught}" + finally: + litellm.suppress_debug_info = previous_suppress_debug_info return _CollectedTrace(tuple(profiler.events), error) @@ -108,6 +124,7 @@ def collect_trace( }, provider_responses=base_fixture.provider_responses, expected_failure=base_fixture.expected_failure, + consume_stream=base_fixture.consume_stream, ) collected: Final = _collect(function, fixture, engine, asynchronous=asynchronous) provider.take_requests(len(fixture.provider_responses)) @@ -129,18 +146,24 @@ def _failure_message(result: tuple[FunctionTraceEvent, ...] | TraceExecutionFail def execute_trace( - route: RouteSpec, scenario: TraceScenario, mode: TraceMode, surface: Surface -) -> TraceComparisonArtifact: - asynchronous: Final = mode == "async" - mappings: Final = scenario.mappings_for(mode) + route: RouteSpec, + scenario: TraceScenario, + surface: Surface, + engine: TraceEngine = "both", +) -> TraceArtifact: + mappings: Final = scenario.mappings scenario_route: Final = RouteSpec( route=route.route, python_entrypoints=route.python_entrypoints, rust_entrypoints=route.rust_entrypoints, fixture=scenario.fixture, ) - python_trace: Final = collect_trace(scenario_route, "python", asynchronous=asynchronous) - rust_trace: Final = collect_trace(scenario_route, "rust", asynchronous=asynchronous) + python_trace: Final = ( + collect_trace(scenario_route, "python", asynchronous=scenario.asynchronous) if engine != "rust" else () + ) + rust_trace: Final = ( + collect_trace(scenario_route, "rust", asynchronous=scenario.asynchronous) if engine != "python" else () + ) python_error: Final = _failure_message(python_trace) rust_error: Final = _failure_message(rust_trace) python_events: Final = python_trace if isinstance(python_trace, tuple) else () @@ -149,28 +172,22 @@ def execute_trace( python: Final = pipeline_projection("python", python_events, mappings) rust: Final = pipeline_projection("rust", rust_events, mappings) except ValueError as error: - return TraceComparisonArtifact.from_traces( + return TraceArtifact.from_traces( + engine=engine, surface=surface, sdk_function=route.route, scenario=scenario.name, - mode=mode, - mappings=mappings, - contract=scenario.contract, python=(), rust=(), - python_unmatched=0, python_error=f"harness: {error}", ) - return TraceComparisonArtifact.from_traces( + return TraceArtifact.from_traces( + engine=engine, surface=surface, sdk_function=route.route, scenario=scenario.name, - mode=mode, - mappings=mappings, - contract=scenario.contract, python=python.steps, rust=rust.steps, - python_unmatched=python.unmatched, python_error=python_error, rust_error=rust_error, ) diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/messages/case.py b/tests/rust-python-harness/strategies/trace_parity/sdk/messages/case.py index 27079c28cd8..211c454eadf 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/messages/case.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/messages/case.py @@ -1,14 +1,26 @@ from __future__ import annotations -import json from typing import Final -from .....shared.parity.recorded_http import HttpHeader, RecordedHttpResponse from .....shared.tracing.steps import Engine, mapping +from ...fixtures import ( + anthropic_response_body, + anthropic_stream_events, + aws_event_stream_response, + json_response, + sse_response, +) from ...models import RouteFixture, RouteSpec, TraceScenario, TraceSuite COMMON_MAPPINGS: Final = ( mapping(rust_span="messages", python_frame=r"anthropic_interface/messages/__init__\.py:\d+ a?create$"), + mapping(span="python_sanitize_empty_content", python_frame=r"strip_empty_content_blocks_from_anthropic_messages$"), + mapping(span="python_sanitize_tool_ids", python_frame=r"sanitize_tool_use_ids_in_anthropic_messages$"), + mapping( + span="python_flatten_web_search", python_frame=r"flatten_unencrypted_web_search_results_in_anthropic_messages$" + ), + mapping(span="python_cache_control", python_frame=r"AnthropicCacheControlHook\.maybe_inject_cache_control$"), + mapping(span="python_pre_request_hooks", python_frame=r"_execute_pre_request_hooks$"), mapping( span="python_messages_provider_config", python_frame=r"ProviderConfigManager\.get_provider_anthropic_messages_config$", @@ -30,10 +42,42 @@ COMMON_MAPPINGS: Final = ( ), mapping(rust_span="http_request", python_frame=r"AsyncHTTPHandler\.post$|HTTPHandler\.post$"), mapping(rust_span="transform_response", python_frame=r"(? RouteFixture: conversation: Final = {"messages": [{"role": "user", "content": "hello"}], "max_tokens": 16} - response: Final = json.dumps( - { - "id": "msg_trace", - "type": "message", - "role": "assistant", - "model": "claude-sonnet-5", - "content": [{"type": "text", "text": "hello"}], - "stop_reason": "end_turn", - "stop_sequence": None, - "usage": {"input_tokens": 2, "output_tokens": 3}, - } - ).encode() return RouteFixture( kwargs={ "model": f"{provider}/claude-sonnet-5", **({"body": {**conversation, "model": "claude-sonnet-5"}} if engine == "rust" else conversation), }, - provider_responses=( - RecordedHttpResponse.from_bytes( - 200, (HttpHeader(name="content-type", value="application/json"),), response - ), - ), + provider_responses=(json_response(anthropic_response_body()),), ) @@ -88,11 +156,170 @@ def _azure_fixture(engine: Engine, _base_url: str) -> RouteFixture: return _fixture(engine, "azure_ai") +def _bedrock_kwargs(engine: Engine) -> dict[str, object]: + conversation: Final = {"messages": [{"role": "user", "content": "hello"}], "max_tokens": 16} + return { + "model": "bedrock/anthropic.claude-3-sonnet-20240229-v1:0", + **( + {"body": {**conversation, "model": "anthropic.claude-3-sonnet-20240229-v1:0"}} + if engine == "rust" + else conversation + ), + "aws_access_key_id": "AKIAIOSFODNN7EXAMPLE", + "aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + "aws_region_name": "us-east-1", + } + + +def _bedrock_fixture(engine: Engine, _base_url: str) -> RouteFixture: + response_fixture: Final = _fixture(engine, "anthropic") + return RouteFixture(kwargs=_bedrock_kwargs(engine), provider_responses=response_fixture.provider_responses) + + +def _bedrock_retry_fixture(engine: Engine, _base_url: str) -> RouteFixture: + success_fixture: Final = _bedrock_fixture(engine, _base_url) + messages: Final = [ + {"role": "user", "content": "hello"}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "old reasoning", "signature": ""}, + {"type": "text", "text": "partial answer"}, + ], + }, + {"role": "user", "content": "continue"}, + ] + kwargs: Final = { + **_bedrock_kwargs(engine), + **( + {"body": {"messages": messages, "max_tokens": 16, "model": "anthropic.claude-3-sonnet-20240229-v1:0"}} + if engine == "rust" + else {"messages": messages} + ), + } + return success_fixture.derive( + kwargs=kwargs, + provider_responses=( + json_response({"message": "messages.1.content.0: Invalid `signature` in `thinking` block"}, status=400), + *success_fixture.provider_responses, + ), + ) + + +def _mock_fixture(engine: Engine, _base_url: str) -> RouteFixture: + fixture: Final = _fixture(engine, "anthropic") + return fixture.derive(kwargs={"mock_response": "hello from mock"}, provider_responses=()) + + +def _provider_error_fixture(engine: Engine, _base_url: str) -> RouteFixture: + fixture: Final = _fixture(engine, "anthropic") + return fixture.derive( + provider_responses=( + json_response( + {"type": "error", "error": {"type": "invalid_request_error", "message": "bad request"}}, + status=400, + ), + ), + expected_failure=True, + ) + + +def _sync_unsupported_fixture(engine: Engine, base_url: str) -> RouteFixture: + if engine == "rust": + return _anthropic_fixture(engine, base_url) + fixture: Final = _fixture(engine, "anthropic") + return fixture.derive(provider_responses=(), expected_failure=True) + + +def _stream_fixture_for(engine: Engine, provider: str) -> RouteFixture: + fixture: Final = _fixture(engine, provider) + return fixture.derive( + kwargs={"stream": True}, + provider_responses=(sse_response(anthropic_stream_events()),), + consume_stream=True, + ) + + +def _stream_fixture(engine: Engine, _base_url: str) -> RouteFixture: + return _stream_fixture_for(engine, "anthropic") + + +def _azure_stream_fixture(engine: Engine, _base_url: str) -> RouteFixture: + return _stream_fixture_for(engine, "azure_ai") + + +def _bedrock_stream_fixture(engine: Engine, base_url: str) -> RouteFixture: + fixture: Final = _bedrock_fixture(engine, base_url) + events: Final = tuple(payload for _, payload in anthropic_stream_events()) + return fixture.derive( + kwargs={"stream": True}, + provider_responses=(aws_event_stream_response(events),), + consume_stream=True, + ) + + +def _bedrock_stream_error_fixture(engine: Engine, base_url: str) -> RouteFixture: + fixture: Final = _bedrock_fixture(engine, base_url) + start: Final = anthropic_stream_events(model="anthropic.claude-3-sonnet-20240229-v1:0")[0][1] + return fixture.derive( + kwargs={"stream": True}, + provider_responses=(aws_event_stream_response((start, {"type": "message_stop"}), corrupt_last_frame=True),), + expected_failure=True, + consume_stream=True, + ) + + SPEC: Final = RouteSpec("messages", ("create", "acreate"), ("messages", "amessages"), _anthropic_fixture) TRACE_SUITE: Final = TraceSuite( route=SPEC, scenarios=( - TraceScenario(name="anthropic", fixture=_anthropic_fixture, mappings=ANTHROPIC_MAPPINGS, modes=("async",)), - TraceScenario(name="azure-ai", fixture=_azure_fixture, mappings=AZURE_MAPPINGS, modes=("async",)), + TraceScenario( + name="async-anthropic", fixture=_anthropic_fixture, mappings=ANTHROPIC_MAPPINGS, asynchronous=True + ), + TraceScenario(name="async-azure-ai", fixture=_azure_fixture, mappings=AZURE_MAPPINGS, asynchronous=True), + TraceScenario(name="async-bedrock", fixture=_bedrock_fixture, mappings=BEDROCK_MAPPINGS, asynchronous=True), + TraceScenario( + name="async-bedrock-invalid-thinking-retry", + fixture=_bedrock_retry_fixture, + mappings=RETRY_MAPPINGS, + asynchronous=True, + ), + TraceScenario(name="async-mock-response", fixture=_mock_fixture, mappings=MOCK_MAPPINGS, asynchronous=True), + TraceScenario( + name="async-anthropic-provider-error", + fixture=_provider_error_fixture, + mappings=ANTHROPIC_FAILURE_MAPPINGS, + asynchronous=True, + ), + TraceScenario( + name="async-anthropic-stream", + fixture=_stream_fixture, + mappings=(*ANTHROPIC_MAPPINGS, *STREAM_MAPPINGS), + asynchronous=True, + ), + TraceScenario( + name="async-azure-ai-stream", + fixture=_azure_stream_fixture, + mappings=(*AZURE_MAPPINGS, *STREAM_MAPPINGS), + asynchronous=True, + ), + TraceScenario( + name="async-bedrock-event-stream", + fixture=_bedrock_stream_fixture, + mappings=(*BEDROCK_MAPPINGS, *STREAM_MAPPINGS), + asynchronous=True, + ), + TraceScenario( + name="async-bedrock-event-stream-error", + fixture=_bedrock_stream_error_fixture, + mappings=(*BEDROCK_MAPPINGS, *STREAM_MAPPINGS, *FAILURE_MAPPINGS), + asynchronous=True, + ), + TraceScenario( + name="sync-unsupported", + fixture=_sync_unsupported_fixture, + mappings=ANTHROPIC_MAPPINGS, + asynchronous=False, + ), ), ) diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py b/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py index 2a4a1b3a152..effd1a0b4f6 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py @@ -358,53 +358,88 @@ TRACE_SUITE: Final = TraceSuite( route=SPEC, scenarios=( TraceScenario( - name="mistral", + name="sync-mistral", fixture=_mistral_fixture, - mappings=COMMON_MAPPINGS, - sync_mappings=(*SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), - async_mappings=(*ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + mappings=(*SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + asynchronous=False, ), TraceScenario( - name="mistral-callback-success", + name="async-mistral", + fixture=_mistral_fixture, + mappings=(*ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + asynchronous=True, + ), + TraceScenario( + name="sync-mistral-callback-success", fixture=_mistral_callback_success_fixture, - mappings=COMMON_MAPPINGS, - sync_mappings=CALLBACK_SUCCESS_SYNC_MAPPINGS, - async_mappings=CALLBACK_SUCCESS_ASYNC_MAPPINGS, + mappings=CALLBACK_SUCCESS_SYNC_MAPPINGS, + asynchronous=False, ), TraceScenario( - name="mistral-callback-failure", + name="async-mistral-callback-success", + fixture=_mistral_callback_success_fixture, + mappings=CALLBACK_SUCCESS_ASYNC_MAPPINGS, + asynchronous=True, + ), + TraceScenario( + name="sync-mistral-callback-failure", fixture=_mistral_callback_failure_fixture, - mappings=(*COMMON_MAPPINGS, FAILURE_CALLBACK_MAPPING), - sync_mappings=CALLBACK_FAILURE_SYNC_MAPPINGS, - async_mappings=CALLBACK_FAILURE_ASYNC_MAPPINGS, + mappings=CALLBACK_FAILURE_SYNC_MAPPINGS, + asynchronous=False, ), TraceScenario( - name="azure-ai", + name="async-mistral-callback-failure", + fixture=_mistral_callback_failure_fixture, + mappings=CALLBACK_FAILURE_ASYNC_MAPPINGS, + asynchronous=True, + ), + TraceScenario( + name="sync-azure-ai", fixture=_azure_fixture, - mappings=AZURE_COMMON_MAPPINGS, - sync_mappings=(*AZURE_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), - async_mappings=(*AZURE_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + mappings=(*AZURE_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + asynchronous=False, ), TraceScenario( - name="azure-document-intelligence", + name="async-azure-ai", + fixture=_azure_fixture, + mappings=(*AZURE_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + asynchronous=True, + ), + TraceScenario( + name="sync-azure-document-intelligence", fixture=_azure_document_intelligence_fixture, - mappings=DOCUMENT_INTELLIGENCE_COMMON_MAPPINGS, - sync_mappings=(*DOCUMENT_INTELLIGENCE_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), - async_mappings=(*DOCUMENT_INTELLIGENCE_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + mappings=(*DOCUMENT_INTELLIGENCE_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + asynchronous=False, ), TraceScenario( - name="vertex-ai", + name="async-azure-document-intelligence", + fixture=_azure_document_intelligence_fixture, + mappings=(*DOCUMENT_INTELLIGENCE_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + asynchronous=True, + ), + TraceScenario( + name="sync-vertex-ai", fixture=_vertex_fixture, - mappings=VERTEX_COMMON_MAPPINGS, - sync_mappings=(*VERTEX_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), - async_mappings=(*VERTEX_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + mappings=(*VERTEX_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + asynchronous=False, ), TraceScenario( - name="vertex-deepseek", + name="async-vertex-ai", + fixture=_vertex_fixture, + mappings=(*VERTEX_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + asynchronous=True, + ), + TraceScenario( + name="sync-vertex-deepseek", fixture=_vertex_deepseek_fixture, - mappings=DEEPSEEK_COMMON_MAPPINGS, - sync_mappings=(*DEEPSEEK_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), - async_mappings=(*DEEPSEEK_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + mappings=(*DEEPSEEK_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + asynchronous=False, + ), + TraceScenario( + name="async-vertex-deepseek", + fixture=_vertex_deepseek_fixture, + mappings=(*DEEPSEEK_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + asynchronous=True, ), ), ) diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/responses/case.py b/tests/rust-python-harness/strategies/trace_parity/sdk/responses/case.py new file mode 100644 index 00000000000..68d2eaba9bf --- /dev/null +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/responses/case.py @@ -0,0 +1,216 @@ +from __future__ import annotations + +from typing import Final + +from .....shared.tracing.steps import Engine, mapping +from ...fixtures import ( + anthropic_response_body, + anthropic_stream_events, + json_response, + responses_body, + responses_stream_events, + sse_response, +) +from ...models import RouteFixture, RouteSpec, TraceScenario, TraceSuite + +COMMON_MAPPINGS: Final = ( + mapping(span="python_responses", python_frame=r"responses/main\.py:\d+ a?responses$"), + mapping( + span="python_responses_provider_config", + python_frame=r"ProviderConfigManager\.get_provider_responses_api_config$", + ), + mapping(rust_span="responses_provider_config"), + mapping(rust_span="validate_environment", python_frame=r"validate_environment$"), + mapping(rust_span="complete_url", python_frame=r"get_complete_url$"), + mapping( + rust_span="transform_request", + python_frame=r"(? RouteFixture: + model: Final = "gpt-5" + return RouteFixture( + kwargs={ + "model": f"{provider}/{model}", + "input": "hello", + **({"body": {"model": model, "input": "hello"}} if engine == "rust" else {}), + }, + provider_responses=(json_response(responses_body(model=model)),), + ) + + +def _openai_fixture(engine: Engine, _base_url: str) -> RouteFixture: + return _native_fixture(engine, "openai") + + +def _azure_fixture(engine: Engine, _base_url: str) -> RouteFixture: + fixture: Final = _native_fixture(engine, "azure") + return fixture.derive(kwargs={"api_version": "2025-04-01-preview"}) + + +def _openai_stream_fixture(engine: Engine, _base_url: str) -> RouteFixture: + fixture: Final = _openai_fixture(engine, _base_url) + return fixture.derive( + kwargs={"stream": True}, + provider_responses=(sse_response(responses_stream_events()),), + consume_stream=True, + ) + + +def _provider_error_fixture(engine: Engine, _base_url: str) -> RouteFixture: + fixture: Final = _openai_fixture(engine, _base_url) + return fixture.derive( + provider_responses=( + json_response({"error": {"message": "bad request", "type": "invalid_request_error"}}, status=400), + ), + expected_failure=True, + ) + + +def _stream_failed_fixture(engine: Engine, base_url: str) -> RouteFixture: + fixture: Final = _openai_fixture(engine, base_url) + failed_response: Final[dict[str, object]] = { + **responses_body(), + "status": "failed", + "output": [], + "error": {"message": "stream failed", "type": "server_error", "code": "server_error"}, + } + events: Final = ( + ( + "response.created", + {"type": "response.created", "response": {**failed_response, "status": "in_progress", "error": None}}, + ), + ("response.failed", {"type": "response.failed", "response": failed_response}), + ) + return fixture.derive( + kwargs={"stream": True}, + provider_responses=(sse_response(events),), + expected_failure=True, + consume_stream=True, + ) + + +def _anthropic_bridge_fixture(engine: Engine, _base_url: str) -> RouteFixture: + return RouteFixture( + kwargs={ + "model": "anthropic/claude-sonnet-5", + "input": "hello", + "max_output_tokens": 16, + **({"body": {"model": "claude-sonnet-5", "input": "hello"}} if engine == "rust" else {}), + }, + provider_responses=(json_response(anthropic_response_body()),), + ) + + +def _anthropic_bridge_stream_fixture(engine: Engine, _base_url: str) -> RouteFixture: + fixture: Final = _anthropic_bridge_fixture(engine, _base_url) + return fixture.derive( + kwargs={"stream": True}, + provider_responses=(sse_response(anthropic_stream_events()),), + consume_stream=True, + ) + + +SPEC: Final = RouteSpec("responses", ("responses", "aresponses"), ("responses", "aresponses"), _openai_fixture) +TRACE_SUITE: Final = TraceSuite( + route=SPEC, + scenarios=( + TraceScenario(name="sync-openai", fixture=_openai_fixture, mappings=COMMON_MAPPINGS, asynchronous=False), + TraceScenario(name="async-openai", fixture=_openai_fixture, mappings=COMMON_MAPPINGS, asynchronous=True), + TraceScenario( + name="sync-openai-stream", + fixture=_openai_stream_fixture, + mappings=(*COMMON_MAPPINGS, *STREAM_MAPPINGS), + asynchronous=False, + ), + TraceScenario( + name="async-openai-stream", + fixture=_openai_stream_fixture, + mappings=(*COMMON_MAPPINGS, *STREAM_MAPPINGS), + asynchronous=True, + ), + TraceScenario( + name="async-openai-provider-error", + fixture=_provider_error_fixture, + mappings=(*COMMON_MAPPINGS, *FAILURE_MAPPINGS), + asynchronous=True, + ), + TraceScenario( + name="async-openai-stream-failed", + fixture=_stream_failed_fixture, + mappings=(*COMMON_MAPPINGS, *STREAM_MAPPINGS, *FAILURE_MAPPINGS), + asynchronous=True, + ), + TraceScenario(name="async-azure", fixture=_azure_fixture, mappings=AZURE_MAPPINGS, asynchronous=True), + TraceScenario( + name="async-anthropic-chat-bridge", + fixture=_anthropic_bridge_fixture, + mappings=BRIDGE_MAPPINGS, + asynchronous=True, + ), + TraceScenario( + name="async-anthropic-chat-bridge-stream", + fixture=_anthropic_bridge_stream_fixture, + mappings=( + *BRIDGE_MAPPINGS, + mapping(span="python_chat_stream_wrapper", python_frame=r"CustomStreamWrapper\.__init__$"), + mapping(span="python_chat_stream_next", python_frame=r"CustomStreamWrapper\.__anext__$"), + mapping( + span="python_responses_bridge_stream_iterator", + python_frame=r"LiteLLMCompletionStreamingIterator\.__init__$|LiteLLMCompletionStreamingIterator\.__anext__$", + ), + ), + asynchronous=True, + ), + ), +) diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/test_core_scenario_matrix.py b/tests/rust-python-harness/strategies/trace_parity/sdk/test_core_scenario_matrix.py new file mode 100644 index 00000000000..d0921398795 --- /dev/null +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/test_core_scenario_matrix.py @@ -0,0 +1,63 @@ +from __future__ import annotations + +from importlib import import_module +from typing import Final, cast + +from ..models import TraceSuite + + +def _suite(module: str) -> TraceSuite: + loaded: Final = import_module(module) + candidate: Final = cast(object, getattr(loaded, "TRACE_SUITE")) + assert isinstance(candidate, TraceSuite) + return candidate + + +def test_core_sdk_scenario_matrix_keeps_distinct_migration_paths() -> None: + chat: Final = _suite("tests.rust-python-harness.strategies.trace_parity.sdk.chat_completions.case") + messages: Final = _suite("tests.rust-python-harness.strategies.trace_parity.sdk.messages.case") + responses: Final = _suite("tests.rust-python-harness.strategies.trace_parity.sdk.responses.case") + + assert {(scenario.name, scenario.asynchronous) for scenario in chat.scenarios} >= { + ("sync-anthropic", False), + ("async-anthropic", True), + ("sync-anthropic-stream", False), + ("async-anthropic-stream", True), + ("async-anthropic-provider-error", True), + ("async-anthropic-stream-error", True), + ("sync-bedrock", False), + ("async-bedrock", True), + ("sync-bedrock-event-stream", False), + ("async-bedrock-event-stream", True), + } + assert {(scenario.name, scenario.asynchronous) for scenario in messages.scenarios} >= { + ("async-anthropic-stream", True), + ("async-azure-ai-stream", True), + ("async-bedrock-event-stream", True), + ("async-bedrock-event-stream-error", True), + ("async-bedrock-invalid-thinking-retry", True), + ("sync-unsupported", False), + } + assert {(scenario.name, scenario.asynchronous) for scenario in responses.scenarios} >= { + ("sync-openai", False), + ("async-openai", True), + ("sync-openai-stream", False), + ("async-openai-stream", True), + ("async-openai-provider-error", True), + ("async-openai-stream-failed", True), + ("async-azure", True), + ("async-anthropic-chat-bridge", True), + ("async-anthropic-chat-bridge-stream", True), + } + + +def test_core_gateway_matrix_keeps_downstream_streams_separate() -> None: + modules: Final = ( + "tests.rust-python-harness.strategies.trace_parity.gateway.chat_completions.case", + "tests.rust-python-harness.strategies.trace_parity.gateway.messages.case", + "tests.rust-python-harness.strategies.trace_parity.gateway.responses.case", + ) + + for module in modules: + suite = _suite(module) + assert any("downstream-stream" in scenario.name for scenario in suite.scenarios) diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/transcription/case.py b/tests/rust-python-harness/strategies/trace_parity/sdk/transcription/case.py index 2fb1054d10a..3b4d2e1447d 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/transcription/case.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/transcription/case.py @@ -92,11 +92,16 @@ TRACE_SUITE: Final = TraceSuite( route=SPEC, scenarios=( TraceScenario( - name="bedrock", + name="sync-bedrock", fixture=_fixture, - mappings=MAPPINGS, - sync_mappings=SYNC_MAPPINGS, - async_mappings=ASYNC_MAPPINGS, + mappings=SYNC_MAPPINGS, + asynchronous=False, + ), + TraceScenario( + name="async-bedrock", + fixture=_fixture, + mappings=ASYNC_MAPPINGS, + asynchronous=True, ), ), ) diff --git a/tests/rust-python-harness/strategies/trace_parity/test_reporting.py b/tests/rust-python-harness/strategies/trace_parity/test_reporting.py index 68264064c92..22cc87592b8 100644 --- a/tests/rust-python-harness/strategies/trace_parity/test_reporting.py +++ b/tests/rust-python-harness/strategies/trace_parity/test_reporting.py @@ -1,58 +1,46 @@ from __future__ import annotations -from collections.abc import Sequence from typing import Final, Literal import pytest from ...shared.reporting.models import CaseResult, Coverage, HarnessCase, ResultArtifact, RunStatus from ...shared.reporting.strategy import ModuleCaseSpec, NotImplementedCaseSpec -from ...shared.tracing.steps import PipelineStep, TraceContract, TraceMapping, mapping +from ...shared.tracing.steps import PipelineStep from . import reporting -from .reporting import TRACE_COMPARISON_ARTIFACT, TraceComparisonArtifact, render_trace_results - -MAPPINGS: Final = ( - mapping(rust_span="ocr", python_frame=r"ocr/main\.py:\d+ a?ocr$"), - mapping(rust_span="http_request", python_frame=r"AsyncHTTPHandler\.post$"), -) +from .reporting import TRACE_ARTIFACT, TraceArtifact, render_trace_results -def _result(comparison: TraceComparisonArtifact) -> CaseResult: +def _result(trace: TraceArtifact) -> CaseResult: case: Final = HarnessCase( strategy_id="trace_parity", strategy_label="Trace parity", - sdk_function=comparison.sdk_function, + sdk_function=trace.sdk_function, spec=ModuleCaseSpec(coverage=Coverage.PARTIAL, module="example"), - surface=comparison.surface, + surface=trace.surface, ) result: Final = CaseResult(case=case) - nodeid: Final = f"trace:sdk:{comparison.sdk_function}:{comparison.scenario}:{comparison.mode}" + nodeid: Final = f"trace:{trace.surface}:{trace.sdk_function}:{trace.scenario}" result.collected.add(nodeid) - result.record( - nodeid, - RunStatus.PASSED, - artifacts=(ResultArtifact(TRACE_COMPARISON_ARTIFACT, comparison.model_dump_json()),), - ) + result.record(nodeid, RunStatus.PASSED, artifacts=(ResultArtifact(TRACE_ARTIFACT, trace.model_dump_json()),)) return result -def _comparison( +def _trace( python: tuple[PipelineStep, ...], rust: tuple[PipelineStep, ...], *, - mappings: Sequence[TraceMapping] = MAPPINGS, rust_error: str | None = None, -) -> TraceComparisonArtifact: - return TraceComparisonArtifact.from_traces( + engine: Literal["python", "rust", "both"] = "both", + scenario: str = "sync-default", +) -> TraceArtifact: + return TraceArtifact.from_traces( + engine=engine, surface="sdk", sdk_function="ocr", - scenario="default", - mode="sync", - mappings=mappings, - contract=TraceContract(), + scenario=scenario, python=python, rust=rust, - python_unmatched=796, rust_error=rust_error, ) @@ -67,107 +55,55 @@ def _events(*items: tuple[str, int, str | None]) -> tuple[PipelineStep, ...]: return tuple(steps) -def test_renderer_shows_matching_python_and_rust_paths() -> None: - rust: Final = _events(("ocr", 0, None), ("http_request", 1, None)) +def test_renderer_prints_python_and_rust_traces_independently() -> None: python: Final = _events( ("ocr", 0, "ocr/main.py:88 aocr"), - ("http_request", 1, "http_handler.py:673 AsyncHTTPHandler.post"), + ("python_prepare", 1, "prep.py:1 python_prepare"), ) - - section: Final = render_trace_results((_result(_comparison(python, rust)),))[0] - report: Final = "\n\n".join(section.blocks) - - assert section.title == "SDK trace comparisons" - assert "Case: ocr" in report - assert "PYTHON (2 steps)\n1 aocr (ocr/main.py:88)\n2 AsyncHTTPHandler.post (http_handler.py:673)" in report - assert "RUST (2 steps)\nocr -> 1 aocr\n http_request -> 2 AsyncHTTPHandler.post" in report - assert "Mapping (identifier -> span)" not in report - assert "Trace: MATCH" in report - assert "Same steps, order, and nesting" in report - assert "Unseen mappings:" not in report - - -def test_renderer_reports_mappings_that_matched_nothing() -> None: - events: Final = _events(("ocr", 0, None)) - - section: Final = render_trace_results((_result(_comparison(events, events)),))[0] - report: Final = "\n\n".join(section.blocks) - - assert "Unseen mappings: http_request" in report - assert "Contract: FAIL" in report - - -def test_renderer_numbers_repeated_span_occurrences() -> None: - mappings: Final = (MAPPINGS[0], MAPPINGS[1]) - rust: Final = _events(("ocr", 0, None), ("http_request", 1, None), ("http_request", 1, None)) - python: Final = _events( - ("ocr", 0, "ocr/main.py:88 aocr"), - ("http_request", 1, "http_handler.py:673 AsyncHTTPHandler.post"), - ("http_request", 1, "http_handler.py:673 AsyncHTTPHandler.post"), - ) - - report: Final = "\n\n".join(render_trace_results((_result(_comparison(python, rust, mappings=mappings)),))[0].blocks) - - assert "http_request#2" in report - - -def test_renderer_accepts_declared_engine_specific_steps() -> None: - mappings: Final = ( - *MAPPINGS[:1], - mapping(span="python_prepare", python_frame=r"python_prepare$"), - mapping(rust_span="rust_prepare"), - ) - python: Final = _events(("ocr", 0, None), ("python_prepare", 1, "prep.py:1 python_prepare")) rust: Final = _events(("ocr", 0, None), ("rust_prepare", 1, None)) - section: Final = render_trace_results((_result(_comparison(python, rust, mappings=mappings)),))[0] + section: Final = render_trace_results((_result(_trace(python, rust)),))[0] report: Final = "\n\n".join(section.blocks) - assert "2 python_prepare (prep.py:1) [python only]" in report - assert "rust_prepare -> [rust only]" in report - assert "Trace: MATCH" in report - assert "Contract: PASS" in report + assert section.title == "SDK traces" + assert "PYTHON (2 steps)\n1 aocr (ocr/main.py:88)\n2 python_prepare (prep.py:1)" in report + assert "RUST (2 steps)\n1 ocr\n2 rust_prepare" in report + assert "python only" not in report + assert "rust only" not in report + assert " -> " not in report + assert "Trace: MATCH" not in report + assert "Trace: DRIFT" not in report + assert "Contract:" not in report -def test_unavailable_check_reports_mode_from_nodeid() -> None: - case: Final = HarnessCase( - strategy_id="trace_parity", - strategy_label="Trace parity", - sdk_function="ocr", - spec=ModuleCaseSpec(coverage=Coverage.PARTIAL, module="example"), - surface="sdk", - ) - result: Final = CaseResult(case=case) - result.collected.add("trace:sdk:ocr:default:sync") - result.record("trace:sdk:ocr:default:sync", RunStatus.ERROR) +@pytest.mark.parametrize( + ("engine", "present", "absent"), + (("python", "PYTHON (1 steps)", "RUST"), ("rust", "RUST (1 steps)", "PYTHON")), +) +def test_renderer_prints_only_selected_engine(engine: Literal["python", "rust"], present: str, absent: str) -> None: + events: Final = _events(("ocr", 0, None)) - section: Final = render_trace_results((result,))[0] - report: Final = "\n\n".join(section.blocks) + report: Final = "\n\n".join(render_trace_results((_result(_trace(events, events, engine=engine)),))[0].blocks) - assert "Case: ocr" in report - assert "Scenario: default / Mode: sync" in report - assert "Trace: NOT AVAILABLE\nTest outcome: error" in report - assert "unknown mode" not in report + assert present in report + assert absent not in report def test_renderer_keeps_collected_trace_when_one_engine_errors() -> None: - python: Final = _events( - ("ocr", 0, "ocr/main.py:88 aocr"), - ("http_request", 1, "http_handler.py:673 AsyncHTTPHandler.post"), + python: Final = _events(("ocr", 0, "ocr/main.py:88 aocr")) + + report: Final = "\n\n".join( + render_trace_results( + (_result(_trace(python, (), rust_error="rust: native Rust bridge must include the trace-parity feature")),) + )[0].blocks ) - section: Final = render_trace_results( - (_result(_comparison(python, (), rust_error="rust: native Rust bridge must include the trace-parity feature")),) - )[0] - report: Final = "\n\n".join(section.blocks) - - assert "PYTHON (2 steps)\n1 aocr (ocr/main.py:88) [python only]" in report + assert "PYTHON (1 steps)\n1 aocr (ocr/main.py:88)" in report assert "Rust error: rust: native Rust bridge must include the trace-parity feature" in report assert "hint: rebuild the native bridge with the trace-parity feature" in report - assert "Contract: FAIL" in report -def test_renderer_groups_all_modes_under_one_case_header() -> None: +def test_unavailable_trace_reports_scenario_from_nodeid() -> None: case: Final = HarnessCase( strategy_id="trace_parity", strategy_label="Trace parity", @@ -176,76 +112,55 @@ def test_renderer_groups_all_modes_under_one_case_header() -> None: surface="sdk", ) result: Final = CaseResult(case=case) - events: Final = _events(("ocr", 0, None)) - modes: Final[tuple[Literal["sync", "async"], ...]] = ("sync", "async") - for mode in modes: - nodeid = f"trace:sdk:ocr:default:{mode}" - result.collected.add(nodeid) - comparison = TraceComparisonArtifact.from_traces( - surface="sdk", - sdk_function="ocr", - scenario="default", - mode=mode, - mappings=MAPPINGS, - contract=TraceContract(), - python=events, - rust=events, - python_unmatched=0, - ) - result.record( - nodeid, - RunStatus.PASSED, - artifacts=(ResultArtifact(TRACE_COMPARISON_ARTIFACT, comparison.model_dump_json()),), - ) + result.collected.add("trace:sdk:ocr:async-error") + result.record("trace:sdk:ocr:async-error", RunStatus.ERROR) - section: Final = render_trace_results((result,))[0] + report: Final = "\n\n".join(render_trace_results((result,))[0].blocks) + + assert "Scenario: async-error" in report + assert "Trace: NOT AVAILABLE\nTest outcome: error" in report + + +def test_renderer_groups_scenarios_under_one_case_header() -> None: + result: Final = _result(_trace(_events(("ocr", 0, None)), (), scenario="sync-default")) + async_trace: Final = _trace((), _events(("ocr", 0, None)), scenario="async-default") + nodeid: Final = "trace:sdk:ocr:async-default" + result.collected.add(nodeid) + result.record(nodeid, RunStatus.PASSED, artifacts=(ResultArtifact(TRACE_ARTIFACT, async_trace.model_dump_json()),)) + + report: Final = render_trace_results((result,))[0].blocks[0] - assert len(section.blocks) == 1 - report: Final = section.blocks[0] assert report.count("Case: ocr") == 1 - assert "Scenario: default / Mode: sync" in report - assert "Scenario: default / Mode: async" in report + assert "Scenario: sync-default" in report + assert "Scenario: async-default" in report def test_renderer_colors_every_trace_line_in_a_terminal(monkeypatch: pytest.MonkeyPatch) -> None: - rust: Final = _events(("ocr", 0, None), ("http_request", 1, None)) - python: Final = _events( - ("ocr", 0, "ocr/main.py:88 aocr"), - ("http_request", 1, "http_handler.py:673 AsyncHTTPHandler.post"), - ) + events: Final = _events(("ocr", 0, "ocr/main.py:88 aocr")) monkeypatch.setattr(reporting.sys.stdout, "isatty", lambda: True) monkeypatch.delenv("NO_COLOR", raising=False) - section: Final = render_trace_results((_result(_comparison(python, rust)),))[0] - report: Final = "\n\n".join(section.blocks) + report: Final = "\n\n".join(render_trace_results((_result(_trace(events, events)),))[0].blocks) - assert "\033[36mPYTHON\033[0m (2 steps)" in report + assert "\033[36mPYTHON\033[0m (1 steps)" in report assert "\033[36m1 aocr (ocr/main.py:88)\033[0m" in report - assert "\033[33mRUST\033[0m (2 steps)" in report - assert "\033[33mocr\033[0m -> \033[36m1 aocr\033[0m" in report - assert "\033[33mhttp_request\033[0m -> \033[36m2 AsyncHTTPHandler.post\033[0m" in report + assert "\033[33mRUST\033[0m (1 steps)" in report + assert "\033[33m1 ocr\033[0m" in report -def test_renderer_groups_cases_and_unavailable_entries_by_surface() -> None: - events: Final = _events(("ocr", 0, None)) - gateway_results: Final = tuple( - CaseResult( - case=HarnessCase( - strategy_id="trace_parity", - strategy_label="Trace parity", - sdk_function=sdk_function, - spec=NotImplementedCaseSpec(reason=f"No {sdk_function} case is registered."), - surface="gateway", - ), - status=RunStatus.NOT_IMPLEMENTED, - ) - for sdk_function in ("ocr", "messages") +def test_renderer_groups_unavailable_entries_by_surface() -> None: + gateway_result: Final = CaseResult( + case=HarnessCase( + strategy_id="trace_parity", + strategy_label="Trace parity", + sdk_function="messages", + spec=NotImplementedCaseSpec(reason="No messages case is registered."), + surface="gateway", + ), + status=RunStatus.NOT_IMPLEMENTED, ) - sections: Final = render_trace_results((_result(_comparison(events, events)), *gateway_results)) + sections: Final = render_trace_results((_result(_trace((), ())), gateway_result)) - assert tuple(section.title for section in sections) == ("SDK trace comparisons", "GATEWAY trace comparisons") - gateway_report: Final = "\n\n".join(sections[1].blocks) - assert gateway_report.count("Not implemented") == 1 - assert "- ocr: No ocr case is registered." in gateway_report - assert "- messages: No messages case is registered." in gateway_report + assert tuple(section.title for section in sections) == ("SDK traces", "GATEWAY traces") + assert "- messages: No messages case is registered." in "\n\n".join(sections[1].blocks) diff --git a/tests/rust-python-harness/strategies/trace_parity/test_runner.py b/tests/rust-python-harness/strategies/trace_parity/test_runner.py index 0fcc5860ff2..4992854e21d 100644 --- a/tests/rust-python-harness/strategies/trace_parity/test_runner.py +++ b/tests/rust-python-harness/strategies/trace_parity/test_runner.py @@ -1,12 +1,20 @@ from __future__ import annotations -from typing import Final +import importlib +from pathlib import Path +from typing import Final, cast + +import pytest + +import litellm from ...shared.reporting.models import Coverage, HarnessCase, HarnessRun, RunStatus, SdkFunction, Surface from ...shared.reporting.strategy import ModuleCaseSpec -from ...shared.tracing.steps import Engine +from ...shared.tracing.steps import Engine, PipelineStep from .models import GatewayRouteSpec, RouteFixture, RouteSpec, TraceScenario, TraceSuite -from .runner import run_trace_mode, scenario_nodeids, validate_trace_suite +from .reporting import TraceArtifact +from .runner import run_trace_cases, run_trace_scenario, runner_selection, scenario_nodeids, validate_trace_suite +from .sdk.execution import execute_trace def _fixture(_engine: Engine, _base_url: str) -> RouteFixture: @@ -27,46 +35,98 @@ def test_scenario_filtering_and_occurrence_node_ids() -> None: suite: Final = TraceSuite( route=RouteSpec("ocr", ("ocr", "aocr"), ("ocr", "aocr"), _fixture), scenarios=( - TraceScenario("one", _fixture, (), modes=("sync", "async")), - TraceScenario("two", _fixture, (), modes=("async",)), + TraceScenario("sync-one", _fixture, (), asynchronous=False), + TraceScenario("async-one", _fixture, (), asynchronous=True), + TraceScenario("async-two", _fixture, (), asynchronous=True), ), ) case: Final = _case() - nodes: Final = scenario_nodeids(suite, case, frozenset({"two"})) + nodes: Final = scenario_nodeids(suite, case, frozenset({"async-two"})) - assert tuple(nodeid for _, _, nodeid in nodes) == ("trace:sdk:ocr:two:async",) + assert tuple(nodeid for _, nodeid in nodes) == ("trace:sdk:ocr:async-two",) + + +def test_python_engine_is_separate_from_scenario_selection() -> None: + assert runner_selection(("mistral", "--engine=python")) == (frozenset({"mistral"}), "python") + + +def test_python_engine_skips_native_bridge(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + runner: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.runner") + case: Final = _case() + selected: list[tuple[frozenset[str], str]] = [] + + def reject_bridge(_repo_root: Path) -> str | None: + raise AssertionError("Python-only tracing must not inspect or build the native bridge") + + def capture_case( + _run: HarnessRun, + _case: HarnessCase, + scenarios: frozenset[str], + _on_update: object, + engine: str, + ) -> None: + selected.append((scenarios, engine)) + + monkeypatch.setattr(runner, "ensure_trace_bridge", reject_bridge) + monkeypatch.setattr(runner, "_run_case", capture_case) + + exit_code, _ = run_trace_cases((case,), tmp_path, lambda _: None, ("mistral", "--engine=python")) + + assert exit_code == 0 + assert selected == [(frozenset({"mistral"}), "python")] + + +def test_expected_provider_failure_omits_feedback_banner( + capsys: pytest.CaptureFixture[str], monkeypatch: pytest.MonkeyPatch +) -> None: + loaded: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.sdk.responses.case") + suite: Final = cast(TraceSuite, loaded.TRACE_SUITE) + scenario: Final = next(item for item in suite.scenarios if item.name == "async-openai-provider-error") + monkeypatch.setattr(litellm, "suppress_debug_info", False) + assert isinstance(suite.route, RouteSpec) + + result: Final = execute_trace(suite.route, scenario, "sdk", engine="python") + + assert result.python_error is None + assert "Give Feedback / Get Help" not in capsys.readouterr().out + assert litellm.suppress_debug_info is False def test_scenario_validation_rejects_duplicate_and_unsafe_names() -> None: route: Final = RouteSpec("ocr", ("ocr", "aocr"), ("ocr", "aocr"), _fixture) duplicate: Final = TraceSuite( route=route, - scenarios=(TraceScenario("same", _fixture, ()), TraceScenario("same", _fixture, ())), + scenarios=( + TraceScenario("sync-same", _fixture, (), asynchronous=False), + TraceScenario("sync-same", _fixture, (), asynchronous=False), + ), + ) + unsafe: Final = TraceSuite( + route=route, scenarios=(TraceScenario("sync-bad:name", _fixture, (), asynchronous=False),) ) - unsafe: Final = TraceSuite(route=route, scenarios=(TraceScenario("bad:name", _fixture, ()),)) case: Final = _case() assert validate_trace_suite(duplicate, case) is not None assert validate_trace_suite(unsafe, case) is not None -def test_scenario_validation_rejects_invalid_modes_and_route_registration() -> None: - invalid_modes: Final = TraceSuite( +def test_scenario_validation_rejects_invalid_names_and_route_registration() -> None: + invalid_name: Final = TraceSuite( route=RouteSpec("ocr", ("ocr", "aocr"), ("ocr", "aocr"), _fixture), - scenarios=(TraceScenario("invalid", _fixture, (), modes=("sync", "sync")),), + scenarios=(TraceScenario("bedrock", _fixture, (), asynchronous=True),), ) wrong_function: Final = TraceSuite( route=RouteSpec("messages", ("create", "acreate"), ("messages", "amessages"), _fixture), - scenarios=(TraceScenario("one", _fixture, ()),), + scenarios=(TraceScenario("sync-one", _fixture, (), asynchronous=False),), ) wrong_surface: Final = TraceSuite( route=GatewayRouteSpec("ocr"), - scenarios=(TraceScenario("one", _fixture, ()),), + scenarios=(TraceScenario("sync-one", _fixture, (), asynchronous=False),), ) case: Final = _case() - assert "unique sync/async modes" in (validate_trace_suite(invalid_modes, case) or "") + assert "start with sync- or async-" in (validate_trace_suite(invalid_name, case) or "") assert "does not match case function" in (validate_trace_suite(wrong_function, case) or "") assert "must use RouteSpec" in (validate_trace_suite(wrong_surface, case) or "") @@ -77,11 +137,35 @@ def test_invalid_route_dispatch_records_harness_error() -> None: result: Final = run.results[case.key] suite: Final = TraceSuite( route=GatewayRouteSpec("ocr"), - scenarios=(TraceScenario("one", _fixture, (), modes=("sync",)),), + scenarios=(TraceScenario("sync-one", _fixture, (), asynchronous=False),), ) - nodeid: Final = "trace:sdk:ocr:one:sync" + nodeid: Final = "trace:sdk:ocr:sync-one" - run_trace_mode(run, result, suite, suite.scenarios[0], "sync", "sdk", nodeid, lambda _: None) + run_trace_scenario(run, result, suite, suite.scenarios[0], "sdk", nodeid, lambda _: None) assert result.outcomes[nodeid] is RunStatus.ERROR assert run.failures == [(nodeid, "gateway route cannot run on the sdk surface")] + + +def test_different_python_and_rust_traces_pass(monkeypatch: pytest.MonkeyPatch) -> None: + runner: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.runner") + case: Final = _case() + run: Final = HarnessRun.from_cases((case,)) + result: Final = run.results[case.key] + suite: Final = TraceSuite( + route=RouteSpec("ocr", ("ocr", "aocr"), ("ocr", "aocr"), _fixture), + scenarios=(TraceScenario("sync-one", _fixture, (), asynchronous=False),), + ) + trace: Final = TraceArtifact.from_traces( + surface="sdk", + sdk_function="ocr", + scenario="sync-one", + python=(PipelineStep(0, None, "python_step", "python.py:1 python_step"),), + rust=(PipelineStep(0, None, "rust_step", "rust_step"),), + ) + monkeypatch.setattr(runner, "_execute_scenario", lambda *_args: trace) + + run_trace_scenario(run, result, suite, suite.scenarios[0], "sdk", "trace:sdk:ocr:sync-one", lambda _: None) + + assert result.outcomes["trace:sdk:ocr:sync-one"] is RunStatus.PASSED + assert run.failures == [] From f486579a2724caa1f5bcf304f5383033ab5828fa Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 14 Sep 2026 11:37:40 -0700 Subject: [PATCH 052/112] fix(harness): rebuild bridge for trace parity --- .../shared/native_build.py | 20 ++++++++++++++++++- .../shared/test_native_build.py | 11 +++++----- 2 files changed, 25 insertions(+), 6 deletions(-) diff --git a/tests/rust-python-harness/shared/native_build.py b/tests/rust-python-harness/shared/native_build.py index b50def98e3a..f67488cecb4 100644 --- a/tests/rust-python-harness/shared/native_build.py +++ b/tests/rust-python-harness/shared/native_build.py @@ -16,6 +16,11 @@ _RUST_ROOT: Final = "litellm-rust" _LOCKFILE: Final = "Cargo.lock" _SOURCE_SUFFIXES: Final = frozenset({".rs", ".toml"}) _FAILURE_OUTPUT_LINES: Final = 15 +_TRACE_CHECK: Final = ( + "from litellm.rust_bridge import get_native_bridge; " + "bridge = get_native_bridge(); " + "raise SystemExit(0 if bridge is not None and getattr(bridge, '_trace', None) is not None else 1)" +) def needs_rebuild(native_mtime: float | None, newest_source_mtime: float | None) -> bool: @@ -73,6 +78,17 @@ def _rebuild(repo_root: Path) -> tuple[bool, str]: return completed.returncode == 0, "\n".join(lines[-_FAILURE_OUTPUT_LINES:]) +def _installed_bridge_has_trace(repo_root: Path) -> bool: + completed: Final = subprocess.run( + (sys.executable, "-c", _TRACE_CHECK), + cwd=repo_root, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + check=False, + ) + return completed.returncode == 0 + + def trace_bridge_error() -> str | None: bridge: Final = get_native_bridge() if bridge is None: @@ -85,7 +101,9 @@ def trace_bridge_error() -> str | None: def ensure_trace_bridge(repo_root: Path) -> str | None: native_path: Final = _native_module_path() native_mtime: Final = native_path.stat().st_mtime if native_path is not None and native_path.exists() else None - rebuild_required: Final = needs_rebuild(native_mtime, _newest_source_mtime(repo_root)) or trace_bridge_error() is not None + rebuild_required: Final = needs_rebuild( + native_mtime, _newest_source_mtime(repo_root) + ) or not _installed_bridge_has_trace(repo_root) if rebuild_required: print(f"Rebuilding native Rust bridge ({BRIDGE_FEATURE} feature)...", flush=True) succeeded: Final diff --git a/tests/rust-python-harness/shared/test_native_build.py b/tests/rust-python-harness/shared/test_native_build.py index 982f787aefa..dc08bc1a2b6 100644 --- a/tests/rust-python-harness/shared/test_native_build.py +++ b/tests/rust-python-harness/shared/test_native_build.py @@ -105,13 +105,14 @@ def test_ensure_trace_bridge_rebuilds_when_trace_feature_is_missing( state.rebuilt = True return True, "" + def fake_get_native_bridge() -> SimpleNamespace: + assert state.rebuilt + return SimpleNamespace(_trace=object()) + monkeypatch.setattr(native_build, "_native_module_path", lambda: native) monkeypatch.setattr(native_build, "_rebuild", fake_rebuild) - monkeypatch.setattr( - native_build, - "get_native_bridge", - lambda: SimpleNamespace(_trace=object() if state.rebuilt else None), - ) + monkeypatch.setattr(native_build, "_installed_bridge_has_trace", lambda repo_root: False) + monkeypatch.setattr(native_build, "get_native_bridge", fake_get_native_bridge) message: Final = native_build.ensure_trace_bridge(tmp_path) From 209913c1c79481a8889d84509e6588891d1b65c9 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 14 Sep 2026 11:38:13 -0700 Subject: [PATCH 053/112] refactor: print rust and python traces independently --- .../unit_tests_mapping/mapping_validator.py | 46 +++++++++---------- 1 file changed, 22 insertions(+), 24 deletions(-) diff --git a/tests/rust-python-harness/strategies/unit_tests_mapping/mapping_validator.py b/tests/rust-python-harness/strategies/unit_tests_mapping/mapping_validator.py index 98ea0b02e68..9dd79e860e6 100644 --- a/tests/rust-python-harness/strategies/unit_tests_mapping/mapping_validator.py +++ b/tests/rust-python-harness/strategies/unit_tests_mapping/mapping_validator.py @@ -38,30 +38,28 @@ def _trace_functions( python_functions: Final[dict[str, PythonFunctionIdentity]] = {} rust_functions: Final[dict[str, RustFunctionIdentity]] = {} for scenario in suite.scenarios: - for mode in scenario.modes: - route: Final = RouteSpec( - route=suite.route.route, - python_entrypoints=suite.route.python_entrypoints, - rust_entrypoints=suite.route.rust_entrypoints, - fixture=scenario.fixture, - ) - python_trace: Final = collect_trace(route, "python", asynchronous=mode == "async") - rust_trace: Final = collect_trace(route, "rust", asynchronous=mode == "async") - if isinstance(python_trace, TraceExecutionFailure): - raise ValueError(f"Python trace discovery failed for {scenario.name}/{mode}: {python_trace.message}") - if isinstance(rust_trace, TraceExecutionFailure): - raise ValueError(f"Rust trace discovery failed for {scenario.name}/{mode}: {rust_trace.message}") - mappings: Final = scenario.mappings_for(mode) - python_projection: Final = pipeline_projection("python", python_trace, mappings) - rust_projection: Final = pipeline_projection("rust", rust_trace, mappings) - for step in python_projection.steps: - if step.span in spec.trace_spans: - function: Final = PythonFunctionIdentity.from_trace(step.raw) - python_functions[function.raw] = function - for step in rust_projection.steps: - if step.span in spec.trace_spans: - function: Final = RustFunctionIdentity.from_trace(step.raw) - rust_functions[step.raw] = function + route: Final = RouteSpec( + route=suite.route.route, + python_entrypoints=suite.route.python_entrypoints, + rust_entrypoints=suite.route.rust_entrypoints, + fixture=scenario.fixture, + ) + python_trace: Final = collect_trace(route, "python", asynchronous=scenario.asynchronous) + rust_trace: Final = collect_trace(route, "rust", asynchronous=scenario.asynchronous) + if isinstance(python_trace, TraceExecutionFailure): + raise ValueError(f"Python trace discovery failed for {scenario.name}: {python_trace.message}") + if isinstance(rust_trace, TraceExecutionFailure): + raise ValueError(f"Rust trace discovery failed for {scenario.name}: {rust_trace.message}") + python_projection: Final = pipeline_projection("python", python_trace, scenario.mappings) + rust_projection: Final = pipeline_projection("rust", rust_trace, scenario.mappings) + for step in python_projection.steps: + if step.span in spec.trace_spans: + function: Final = PythonFunctionIdentity.from_trace(step.raw) + python_functions[function.raw] = function + for step in rust_projection.steps: + if step.span in spec.trace_spans: + function: Final = RustFunctionIdentity.from_trace(step.raw) + rust_functions[step.raw] = function if not python_functions or not rust_functions: raise ValueError(f"Python trace discovery found no functions for spans: {', '.join(spec.trace_spans)}") return ( From 32699a0be91c66c907e93c61eb2969d3c52dfee2 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 13 Sep 2026 21:46:25 -0700 Subject: [PATCH 054/112] build(deps): re-suppress GHSA-h7x2-h6g9-p789 in osv-scan, mlflow still has no fixed release (#41036) Co-authored-by: mateo Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- osv-scanner.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/osv-scanner.toml b/osv-scanner.toml index 3e070fc8cf7..482254d4da6 100644 --- a/osv-scanner.toml +++ b/osv-scanner.toml @@ -5,5 +5,5 @@ reason = "diskcache has no fixed release published; remove this entry once one e [[IgnoredVulns]] id = "GHSA-h7x2-h6g9-p789" -ignoreUntil = 2026-09-14 -reason = "mlflow has no fixed release published; remove this entry once one exists" +ignoreUntil = 2026-10-14 +reason = "mlflow has no fixed release published (3.16.0, 2026-09-04, and master still store gateway secret api_base unvalidated); remove this entry once one exists" From ad5b87eb9149de28b39cd1b1980fa093f408d673 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 14 Sep 2026 11:19:28 -0700 Subject: [PATCH 055/112] fix: capture reused worker threads in Python traces --- .../shared/tracing/profiler.py | 68 +++++++---- .../shared/tracing/test_profiler.py | 113 ++++++++++++++++-- 2 files changed, 150 insertions(+), 31 deletions(-) diff --git a/tests/rust-python-harness/shared/tracing/profiler.py b/tests/rust-python-harness/shared/tracing/profiler.py index abfb6a2425d..e72a36a5289 100644 --- a/tests/rust-python-harness/shared/tracing/profiler.py +++ b/tests/rust-python-harness/shared/tracing/profiler.py @@ -2,7 +2,8 @@ from __future__ import annotations import sys import threading -from collections.abc import Generator, Iterator, Mapping +import warnings +from collections.abc import Callable, Generator, Iterator, Mapping from contextlib import contextmanager from dataclasses import dataclass from functools import lru_cache @@ -32,6 +33,7 @@ class PythonProfiler: self._source_root: Final = str(source_root.resolve()) + "/" self._seen_frames: Final[set[FrameType]] = set() self._event_ids: Final[dict[FrameType, int]] = {} + self._lock: Final = threading.Lock() self.events: Final[list[FunctionTraceEvent]] = [] def __call__(self, frame: FrameType, event: str, _arg: object) -> None: @@ -40,14 +42,15 @@ class PythonProfiler: function_name: Final = self.function_name(frame) if function_name is None: return - event_id: Final = len(self.events) - parent_id: Final = next( - (self._event_ids[ancestor] for ancestor in _frame_ancestors(frame) if ancestor in self._event_ids), - None, - ) - self._seen_frames.add(frame) - self._event_ids[frame] = event_id - self.events.append(FunctionTraceEvent(id=event_id, parent_id=parent_id, function=function_name)) + with self._lock: + event_id: Final = len(self.events) + parent_id: Final = next( + (self._event_ids[ancestor] for ancestor in _frame_ancestors(frame) if ancestor in self._event_ids), + None, + ) + self._seen_frames.add(frame) + self._event_ids[frame] = event_id + self.events.append(FunctionTraceEvent(id=event_id, parent_id=parent_id, function=function_name)) def function_name(self, frame: FrameType) -> str | None: code: Final = frame.f_code @@ -137,21 +140,51 @@ def _frame_ancestors(frame: FrameType) -> Generator[FrameType]: @contextmanager -def profile_python(source_root: Path, *, threads: bool = False) -> Generator[PythonProfiler]: - profiler: Final = PythonProfiler(source_root) +def _installed_profiler(profiler: Callable[[FrameType, str, object], None], *, threads: bool) -> Generator[None]: + if threads and sys.version_info >= (3, 12): + tool_id: Final = next((slot for slot in (2, 3, 4, 0, 1, 5) if sys.monitoring.get_tool(slot) is None), None) + if tool_id is None: + raise RuntimeError("no sys.monitoring tool ID is available for Python trace collection") + + def started(_code: CodeType, _offset: int) -> None: + profiler(sys._getframe(1), "call", None) + + sys.monitoring.use_tool_id(tool_id, "litellm-python-trace") + try: + sys.monitoring.register_callback(tool_id, sys.monitoring.events.PY_START, started) + sys.monitoring.set_events(tool_id, sys.monitoring.events.PY_START) + yield + finally: + sys.monitoring.set_events(tool_id, 0) + sys.monitoring.register_callback(tool_id, sys.monitoring.events.PY_START, None) + sys.monitoring.free_tool_id(tool_id) + return + if threads: + warnings.warn( + "Python <3.12 cannot trace existing worker threads; use Python 3.12+ for complete threaded traces", + RuntimeWarning, + stacklevel=3, + ) previous_thread: Final = threading.getprofile() if threads: threading.setprofile(profiler) previous: Final = sys.getprofile() sys.setprofile(profiler) try: - yield profiler + yield finally: sys.setprofile(previous) if threads: threading.setprofile(previous_thread) +@contextmanager +def profile_python(source_root: Path, *, threads: bool = False) -> Generator[PythonProfiler]: + profiler: Final = PythonProfiler(source_root) + with _installed_profiler(profiler, threads=threads): + yield profiler + + @contextmanager def profile_python_function_usage( source_root: Path, @@ -160,14 +193,5 @@ def profile_python_function_usage( threads: bool = False, ) -> Generator[PythonFunctionUsageProfiler]: profiler: Final = PythonFunctionUsageProfiler(source_root, functions) - previous_thread: Final = threading.getprofile() - if threads: - threading.setprofile(profiler) - previous: Final = sys.getprofile() - sys.setprofile(profiler) - try: + with _installed_profiler(profiler, threads=threads): yield profiler - finally: - sys.setprofile(previous) - if threads: - threading.setprofile(previous_thread) diff --git a/tests/rust-python-harness/shared/tracing/test_profiler.py b/tests/rust-python-harness/shared/tracing/test_profiler.py index 616e9c23e75..b9a07f249a3 100644 --- a/tests/rust-python-harness/shared/tracing/test_profiler.py +++ b/tests/rust-python-harness/shared/tracing/test_profiler.py @@ -4,9 +4,10 @@ import asyncio import sys import threading from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor from functools import wraps from pathlib import Path -from types import FunctionType +from types import FrameType, FunctionType from typing import Final, ParamSpec, TypeVar, cast import pytest @@ -41,11 +42,12 @@ def _events_named(profiler: PythonProfiler, name: str) -> tuple[FunctionTraceEve return tuple(event for event in profiler.events if event.function.endswith(name)) -def test_profiler_keeps_repeated_calls() -> None: +@pytest.mark.parametrize("threads", (False, True)) +def test_profiler_keeps_repeated_calls(threads: bool) -> None: def called() -> None: return None - with profile_python(Path(__file__).parent) as profiler: + with profile_python(Path(__file__).parent, threads=threads) as profiler: called() called() @@ -60,14 +62,15 @@ def test_profiler_qualifies_decorated_methods_by_class() -> None: assert _module_qualnames(__name__)[cast(FunctionType, Decorated.call.__wrapped__).__code__] == "Decorated.call" -def test_profiler_records_real_frame_ancestry() -> None: +@pytest.mark.parametrize("threads", (False, True)) +def test_profiler_records_real_frame_ancestry(threads: bool) -> None: def called() -> None: return None def outer() -> None: called() - with profile_python(Path(__file__).parent) as profiler: + with profile_python(Path(__file__).parent, threads=threads) as profiler: outer() outer_event, called_event = (event for event in profiler.events if event.function.endswith(("outer", "called"))) @@ -84,18 +87,20 @@ def test_profiler_restores_previous_profiler_after_failure() -> None: assert sys.getprofile() is previous -def test_profiler_does_not_count_coroutine_resumption_as_another_call() -> None: +@pytest.mark.parametrize("threads", (False, True)) +def test_profiler_does_not_count_coroutine_resumption_as_another_call(threads: bool) -> None: async def suspended() -> None: await asyncio.sleep(0) await asyncio.sleep(0) - with profile_python(Path(__file__).parent) as profiler: + with profile_python(Path(__file__).parent, threads=threads) as profiler: asyncio.run(suspended()) assert len(_events_named(profiler, "suspended")) == 1 -def test_profiler_preserves_parent_across_coroutine_suspension() -> None: +@pytest.mark.parametrize("threads", (False, True)) +def test_profiler_preserves_parent_across_coroutine_suspension(threads: bool) -> None: def called() -> None: return None @@ -103,7 +108,7 @@ def test_profiler_preserves_parent_across_coroutine_suspension() -> None: await asyncio.sleep(0) called() - with profile_python(Path(__file__).parent) as profiler: + with profile_python(Path(__file__).parent, threads=threads) as profiler: asyncio.run(suspended()) suspended_event: Final = _events_named(profiler, "suspended")[0] @@ -124,6 +129,96 @@ def test_profiler_captures_worker_threads_when_enabled() -> None: assert called_event.parent_id is None +@pytest.mark.skipif(sys.version_info < (3, 12), reason="existing worker capture requires sys.monitoring") +@pytest.mark.parametrize("prewarm", (False, True)) +def test_profiler_captures_reused_workers_without_leaking_between_sessions(prewarm: bool) -> None: + def called() -> None: + return None + + with ThreadPoolExecutor(max_workers=1) as executor: + if prewarm: + executor.submit(called).result(timeout=5) + with profile_python(Path(__file__).parent, threads=True) as first: + executor.submit(called).result(timeout=5) + executor.submit(called).result(timeout=5) + with profile_python(Path(__file__).parent, threads=True) as second: + executor.submit(called).result(timeout=5) + executor.submit(called).result(timeout=5) + + assert len(_events_named(first, "called")) == 1 + assert len(_events_named(second, "called")) == 1 + + +def test_profiler_restores_main_and_worker_hooks_after_failure() -> None: + previous: Final = sys.getprofile() + previous_thread: Final = threading.getprofile() + + with ThreadPoolExecutor(max_workers=1) as executor: + worker_previous: Final = executor.submit(sys.getprofile).result(timeout=5) + with pytest.raises(RuntimeError, match="stop"): + with profile_python(Path(__file__).parent, threads=True): + raise RuntimeError("stop") + assert executor.submit(sys.getprofile).result(timeout=5) is worker_previous + + assert sys.getprofile() is previous + assert threading.getprofile() is previous_thread + + +@pytest.mark.skipif(sys.version_info < (3, 12), reason="existing worker capture requires sys.monitoring") +def test_function_usage_profiler_captures_reused_workers() -> None: + def selected() -> None: + return None + + function: Final = f"{Path(__file__).name}:{selected.__code__.co_firstlineno} {selected.__qualname__}" + with ThreadPoolExecutor(max_workers=1) as executor: + executor.submit(selected).result(timeout=5) + with profile_python_function_usage(Path(__file__).parent, frozenset((function,)), threads=True) as profiler: + executor.submit(selected).result(timeout=5) + + assert profiler.called == {function} + + +@pytest.mark.skipif(sys.version_info < (3, 12), reason="independent thread hooks require sys.monitoring") +def test_threaded_profiler_preserves_custom_worker_hook_and_releases_monitoring_slot() -> None: + def worker_hook(_frame: FrameType, _event: str, _arg: object) -> None: + return None + + def fail_with_profile(executor: ThreadPoolExecutor) -> None: + with profile_python(Path(__file__).parent, threads=True): + assert executor.submit(sys.getprofile).result(timeout=5) is worker_hook + raise RuntimeError("stop") + + tools_before: Final = tuple(sys.monitoring.get_tool(slot) for slot in range(6)) + with ThreadPoolExecutor(max_workers=1, initializer=lambda: sys.setprofile(worker_hook)) as executor: + assert executor.submit(sys.getprofile).result(timeout=5) is worker_hook + with pytest.raises(RuntimeError, match="stop"): + fail_with_profile(executor) + assert executor.submit(sys.getprofile).result(timeout=5) is worker_hook + + assert tuple(sys.monitoring.get_tool(slot) for slot in range(6)) == tools_before + + +@pytest.mark.skipif(sys.version_info < (3, 12), reason="existing worker capture requires sys.monitoring") +def test_threaded_profiler_keeps_concurrent_event_ids_and_parent_links() -> None: + def child() -> None: + return None + + def parent() -> None: + child() + + with ThreadPoolExecutor(max_workers=4) as executor: + with profile_python(Path(__file__).parent, threads=True) as profiler: + futures: Final = tuple(executor.submit(parent) for _ in range(200)) + for future in futures: + future.result(timeout=5) + + parent_ids: Final = frozenset(event.id for event in _events_named(profiler, "parent")) + children: Final = _events_named(profiler, "child") + assert len(parent_ids) == len(children) == 200 + assert frozenset(event.parent_id for event in children) == parent_ids + assert tuple(event.id for event in profiler.events) == tuple(range(len(profiler.events))) + + def test_function_usage_profiler_records_only_selected_functions() -> None: def selected() -> None: return None From 561533c596f8d2ac41ea6bd5bd166c1ae9ca19af Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 14 Sep 2026 11:33:10 -0700 Subject: [PATCH 056/112] wip --- .../src/bin/trace_parity_gateway.rs | 4 +- .../crates/ai-gateway/src/trace_parity.rs | 10 +- tests/rust-python-harness/AGENTS.md | 2 +- tests/rust-python-harness/cli/__init__.py | 15 +- tests/rust-python-harness/cli/catalog.py | 19 +- tests/rust-python-harness/cli/commands.py | 6 +- tests/rust-python-harness/cli/test_cli.py | 21 +- tests/rust-python-harness/conftest.py | 4 +- .../shared/native_build.py | 3 +- .../shared/parity/fixtures/recording.py | 1 + .../shared/parity/fixtures/store.py | 2 + .../shared/parity/fixtures/test_pipeline.py | 3 +- .../shared/parity/fixtures/test_recording.py | 1 + .../shared/reporting/models.py | 8 +- .../shared/reporting/strategy.py | 8 + .../shared/test_native_build.py | 16 +- .../shared/tracing/steps.py | 11 +- .../shared/tracing/test_steps.py | 4 +- .../e2e_parity/sdk/ocr/fixtures/reducto.py | 4 +- .../e2e_parity/sdk/ocr/test_sdk_parity.py | 5 +- .../strategies/trace_parity/AGENTS.md | 2 +- .../strategies/trace_parity/__init__.py | 34 ++- .../strategies/trace_parity/fixtures.py | 176 ++++++++++++ .../gateway/chat_completions/case.py | 65 +++++ .../trace_parity/gateway/execution.py | 47 +-- .../trace_parity/gateway/messages/case.py | 48 ++-- .../trace_parity/gateway/responses/case.py | 63 ++++ .../strategies/trace_parity/models.py | 44 ++- .../strategies/trace_parity/reporting.py | 225 +++------------ .../strategies/trace_parity/runner.py | 90 +++--- .../trace_parity/sdk/chat_completions/case.py | 172 ++++++++--- .../strategies/trace_parity/sdk/execution.py | 81 +++--- .../trace_parity/sdk/messages/case.py | 269 ++++++++++++++++-- .../strategies/trace_parity/sdk/ocr/case.py | 91 ++++-- .../trace_parity/sdk/responses/case.py | 216 ++++++++++++++ .../sdk/test_core_scenario_matrix.py | 63 ++++ .../trace_parity/sdk/transcription/case.py | 13 +- .../strategies/trace_parity/test_reporting.py | 243 +++++----------- .../strategies/trace_parity/test_runner.py | 120 ++++++-- 39 files changed, 1556 insertions(+), 653 deletions(-) create mode 100644 tests/rust-python-harness/strategies/trace_parity/fixtures.py create mode 100644 tests/rust-python-harness/strategies/trace_parity/gateway/chat_completions/case.py create mode 100644 tests/rust-python-harness/strategies/trace_parity/gateway/responses/case.py create mode 100644 tests/rust-python-harness/strategies/trace_parity/sdk/responses/case.py create mode 100644 tests/rust-python-harness/strategies/trace_parity/sdk/test_core_scenario_matrix.py diff --git a/litellm-rust/crates/ai-gateway/src/bin/trace_parity_gateway.rs b/litellm-rust/crates/ai-gateway/src/bin/trace_parity_gateway.rs index 9036deb9871..e247c650fad 100644 --- a/litellm-rust/crates/ai-gateway/src/bin/trace_parity_gateway.rs +++ b/litellm-rust/crates/ai-gateway/src/bin/trace_parity_gateway.rs @@ -5,6 +5,7 @@ use serde_json::Value; #[derive(Deserialize)] struct Input { + path: String, model_alias: String, provider_model: String, api_base: String, @@ -21,7 +22,8 @@ async fn main() { Ok(input) => input, Err(error) => fail(error), }; - let result = litellm_ai_gateway::trace_parity::traced_messages_request( + let result = litellm_ai_gateway::trace_parity::traced_request( + input.path, input.model_alias, input.provider_model, input.api_base, diff --git a/litellm-rust/crates/ai-gateway/src/trace_parity.rs b/litellm-rust/crates/ai-gateway/src/trace_parity.rs index 00c9b53e691..7540a71fb12 100644 --- a/litellm-rust/crates/ai-gateway/src/trace_parity.rs +++ b/litellm-rust/crates/ai-gateway/src/trace_parity.rs @@ -29,14 +29,15 @@ pub struct TracedGatewayResponse { pub trace: Vec, } -pub async fn traced_messages_request( +pub async fn traced_request( + path: String, model_alias: String, provider_model: String, api_base: String, body: Value, ) -> TracedGatewayResponse { let trace = litellm_core::observability::FunctionTrace::default(); - let result = messages_request(model_alias, provider_model, api_base, body) + let result = request(path, model_alias, provider_model, api_base, body) .with_subscriber(trace.dispatcher()) .await; let events = trace.events(); @@ -54,7 +55,8 @@ pub async fn traced_messages_request( } } -pub async fn messages_request( +pub async fn request( + path: String, model_alias: String, provider_model: String, api_base: String, @@ -75,7 +77,7 @@ pub async fn messages_request( }; let request = Request::builder() .method("POST") - .uri("/v1/messages") + .uri(path) .header(AUTHORIZATION, "Bearer trace-master-key") .header(CONTENT_TYPE, "application/json") .body(Body::from(body.to_string())) diff --git a/tests/rust-python-harness/AGENTS.md b/tests/rust-python-harness/AGENTS.md index 017668d4289..ec973035785 100644 --- a/tests/rust-python-harness/AGENTS.md +++ b/tests/rust-python-harness/AGENTS.md @@ -63,7 +63,7 @@ tests/rust-python-harness/ - Examples: `run e2e_parity --surface sdk --function ocr`, `run unit_tests_parity --function ocr --pytest-arg=-x`, or `run all --function ocr` - `cli/catalog.py` discovers strategies, validates their Python definitions, and orders them; `cli/__init__.py` builds the Click command tree; `cli/commands.py` runs selected cases - `e2e_parity/` compares SDK objects, exceptions, callbacks, and streams, or gateway HTTP responses -- `trace_parity/` compares mapped operations, call counts, and required execution ordering; before running it rebuilds the native bridge with the `trace-parity` feature whenever `litellm-rust` sources are newer than the installed extension (`shared/native_build.py`) +- `trace_parity/` prints filtered Python and Rust execution traces without comparing them; before running it rebuilds the native bridge with the `trace-parity` feature whenever `litellm-rust` sources are newer than the installed extension (`shared/native_build.py`) - E2E and trace strategies load their registered module cases and run surface-specific execution from their folders - `unit_tests_mapping/contracts.py` owns typed harness-side mapping contracts, per-function contracts live below `cases/`, and `mappings.py` exports the registry; live test discovery derives unmapped Python and Rust-only tests without an exhaustive manifest - `unit_tests_mapping/runner.py` validates confirmed mappings against the live Python and Rust inventories and attaches the derived status report diff --git a/tests/rust-python-harness/cli/__init__.py b/tests/rust-python-harness/cli/__init__.py index 13b995825dd..d2bfdc55b19 100644 --- a/tests/rust-python-harness/cli/__init__.py +++ b/tests/rust-python-harness/cli/__init__.py @@ -58,16 +58,29 @@ def _strategy_command(strategy: Strategy) -> click.Command: help=runner_argument.help, ) ) + for runner_option in strategy.definition.runner_options: + name: Final = runner_option.option.removeprefix("--").replace("-", "_") + params.append( + click.Option( + (runner_option.option, name), + type=click.Choice(runner_option.choices), + help=runner_option.help, + ) + ) def run_strategy( sdk_functions: tuple[str, ...], surface: str | None = None, runner_args: tuple[str, ...] = (), + **runner_options: str | None, ) -> int: selected_functions: Final = cast(frozenset[SdkFunction], frozenset(sdk_functions)) selected_surface: Final = cast(Surface | None, surface) cases: Final = select_cases((strategy,), selected_functions, selected_surface) - return run_command((strategy,), cases, runner_args) + option_args: Final = tuple( + f"--{name.replace('_', '-')}={value}" for name, value in runner_options.items() if value is not None + ) + return run_command((strategy,), cases, (*runner_args, *option_args)) return click.Command( strategy.id, diff --git a/tests/rust-python-harness/cli/catalog.py b/tests/rust-python-harness/cli/catalog.py index 03eb032d9c6..073873d0121 100644 --- a/tests/rust-python-harness/cli/catalog.py +++ b/tests/rust-python-harness/cli/catalog.py @@ -21,9 +21,7 @@ def _load_strategy_module(name: str, folder: Path, prefix: str | None) -> Module if prefix is not None: return importlib.import_module(f"{prefix}.{name}") module_name: Final = _synthetic_module_name(folder) - spec: Final = importlib.util.spec_from_file_location( - module_name, folder / "__init__.py" - ) + spec: Final = importlib.util.spec_from_file_location(module_name, folder / "__init__.py") if spec is None or spec.loader is None: raise ValueError(f"{folder}: cannot load strategy package") module: Final = importlib.util.module_from_spec(spec) @@ -59,9 +57,7 @@ def _load_strategy(name: str, folder: Path, prefix: str | None) -> Strategy: if duplicates: raise ValueError(f"{folder}: duplicate strategy cases: {duplicates}") expected: Final = frozenset( - (surface, function) - for surface in (definition.surfaces or (None,)) - for function in SDK_FUNCTIONS + (surface, function) for surface in (definition.surfaces or (None,)) for function in SDK_FUNCTIONS ) actual: Final = frozenset(keys) if actual != expected: @@ -73,8 +69,7 @@ def _load_strategy(name: str, folder: Path, prefix: str | None) -> Strategy: incompatible: Final = tuple( (case.surface, case.sdk_function) for case in definition.cases - if case.spec.disposition is CaseDisposition.RUNNABLE - and not isinstance(case.spec, definition.runnable_spec) + if case.spec.disposition is CaseDisposition.RUNNABLE and not isinstance(case.spec, definition.runnable_spec) ) if incompatible: raise ValueError(f"{folder}: runnable cases do not match {definition.runnable_spec.__name__}: {incompatible}") @@ -102,14 +97,10 @@ def _load_strategy(name: str, folder: Path, prefix: str | None) -> Strategy: def load_catalog(root: Path | None = None) -> tuple[Strategy, ...]: resolved: Final = STRATEGIES_ROOT if root is None else root prefix: Final = _STRATEGIES_PACKAGE.__name__ if resolved == STRATEGIES_ROOT else None - folders: Final = tuple( - info.name for info in pkgutil.iter_modules([str(resolved)]) if info.ispkg - ) + folders: Final = tuple(info.name for info in pkgutil.iter_modules([str(resolved)]) if info.ispkg) if not folders: raise ValueError(f"No strategy packages found below {resolved}") - strategies: Final = tuple( - _load_strategy(name, resolved / name, prefix) for name in sorted(folders) - ) + strategies: Final = tuple(_load_strategy(name, resolved / name, prefix) for name in sorted(folders)) ids: Final = [strategy.id for strategy in strategies] if len(set(ids)) != len(ids): raise ValueError(f"Duplicate strategy id in {resolved}") diff --git a/tests/rust-python-harness/cli/commands.py b/tests/rust-python-harness/cli/commands.py index f94c3277dc2..e51bbb5d966 100644 --- a/tests/rust-python-harness/cli/commands.py +++ b/tests/rust-python-harness/cli/commands.py @@ -21,8 +21,7 @@ def select_cases( case for strategy in strategies for case in strategy.cases - if (not sdk_functions or case.sdk_function in sdk_functions) - and (surface is None or case.surface == surface) + if (not sdk_functions or case.sdk_function in sdk_functions) and (surface is None or case.surface == surface) ) @@ -32,8 +31,7 @@ def run_command( runner_args: Sequence[str] = (), ) -> int: grouped: Final = { - strategy.id: tuple(case for case in cases if case.strategy_id == strategy.id) - for strategy in strategies + strategy.id: tuple(case for case in cases if case.strategy_id == strategy.id) for strategy in strategies } visible: Final = tuple(strategy for strategy in strategies if grouped[strategy.id]) runners: Final = tuple(replace(strategy, cases=grouped[strategy.id]) for strategy in visible) diff --git a/tests/rust-python-harness/cli/test_cli.py b/tests/rust-python-harness/cli/test_cli.py index 226c89843d0..5641aa8a539 100644 --- a/tests/rust-python-harness/cli/test_cli.py +++ b/tests/rust-python-harness/cli/test_cli.py @@ -242,7 +242,7 @@ def _assert_unavailable_cell(strategy: Strategy, case: HarnessCase, section_titl def test_every_unavailable_case_finishes_and_explains_itself() -> None: section_titles: Final = { "e2e_parity": "End-to-end parity outcomes", - "trace_parity": "trace comparisons", + "trace_parity": "traces", "unit_tests_mapping": "Python/Rust unit-test mappings", "unit_tests_parity": "Python backend parity outcomes", "unit_tests_rust": "Native Rust unit-test outcomes", @@ -359,6 +359,25 @@ def test_strategy_command_forwards_repeated_filters_and_runner_arguments( ] +def test_trace_command_forwards_engine_and_scenario(monkeypatch: pytest.MonkeyPatch) -> None: + cli: Final = importlib.import_module("tests.rust-python-harness.cli") + captured: list[tuple[str, ...]] = [] + + def capture_run( + strategies: Sequence[Strategy], + cases: Sequence[HarnessCase], + runner_args: Sequence[str] = (), + ) -> int: + del strategies, cases + captured.append(tuple(runner_args)) + return 0 + + monkeypatch.setattr(cli, "run_command", capture_run) + + assert main(["run", "trace_parity", "--scenario", "async-mistral", "--engine", "python"]) == 0 + assert captured == [("async-mistral", "--engine=python")] + + def test_omitted_surface_selects_every_strategy_surface(monkeypatch: pytest.MonkeyPatch) -> None: cli: Final = importlib.import_module("tests.rust-python-harness.cli") selected: list[str] = [] diff --git a/tests/rust-python-harness/conftest.py b/tests/rust-python-harness/conftest.py index d50d0fa4204..c487b25a9cc 100644 --- a/tests/rust-python-harness/conftest.py +++ b/tests/rust-python-harness/conftest.py @@ -19,9 +19,7 @@ def subprocess_test_environment(monkeypatch: pytest.MonkeyPatch) -> None: def cargo_project(tmp_path: Path) -> Callable[[str, str], Path]: def create(package: str, source: str) -> Path: manifest: Final = tmp_path / "Cargo.toml" - manifest.write_text( - f'[package]\nname = "{package}"\nversion = "0.1.0"\nedition = "2021"\n[workspace]\n' - ) + manifest.write_text(f'[package]\nname = "{package}"\nversion = "0.1.0"\nedition = "2021"\n[workspace]\n') (tmp_path / "src").mkdir() (tmp_path / "src/lib.rs").write_text(source) return manifest diff --git a/tests/rust-python-harness/shared/native_build.py b/tests/rust-python-harness/shared/native_build.py index 8693cf3bac2..b50def98e3a 100644 --- a/tests/rust-python-harness/shared/native_build.py +++ b/tests/rust-python-harness/shared/native_build.py @@ -85,7 +85,8 @@ def trace_bridge_error() -> str | None: def ensure_trace_bridge(repo_root: Path) -> str | None: native_path: Final = _native_module_path() native_mtime: Final = native_path.stat().st_mtime if native_path is not None and native_path.exists() else None - if needs_rebuild(native_mtime, _newest_source_mtime(repo_root)): + rebuild_required: Final = needs_rebuild(native_mtime, _newest_source_mtime(repo_root)) or trace_bridge_error() is not None + if rebuild_required: print(f"Rebuilding native Rust bridge ({BRIDGE_FEATURE} feature)...", flush=True) succeeded: Final output: Final diff --git a/tests/rust-python-harness/shared/parity/fixtures/recording.py b/tests/rust-python-harness/shared/parity/fixtures/recording.py index ee0d4b6de7b..cb6ed968b90 100644 --- a/tests/rust-python-harness/shared/parity/fixtures/recording.py +++ b/tests/rust-python-harness/shared/parity/fixtures/recording.py @@ -191,6 +191,7 @@ class _RecordingHandler(LocalHttpHandler): self.end_headers() self.wfile.write(body) + def _recording_provider(spec: UpstreamEndpoint) -> AbstractContextManager[_RecordingProvider]: return serve_in_thread(_RecordingProvider(spec)) diff --git a/tests/rust-python-harness/shared/parity/fixtures/store.py b/tests/rust-python-harness/shared/parity/fixtures/store.py index 270af2a7625..1145b5d7d27 100644 --- a/tests/rust-python-harness/shared/parity/fixtures/store.py +++ b/tests/rust-python-harness/shared/parity/fixtures/store.py @@ -15,6 +15,8 @@ from .cassette import deserialize_cassette, serialize_cassette from .recording import RecordedInteraction FIXTURE_SCHEMA_VERSION: Final = 1 + + class FixtureInput(Protocol): def canonical_input(self) -> dict[str, object]: ... diff --git a/tests/rust-python-harness/shared/parity/fixtures/test_pipeline.py b/tests/rust-python-harness/shared/parity/fixtures/test_pipeline.py index 4535ba05bf6..b162e4949d2 100644 --- a/tests/rust-python-harness/shared/parity/fixtures/test_pipeline.py +++ b/tests/rust-python-harness/shared/parity/fixtures/test_pipeline.py @@ -45,8 +45,8 @@ class _Upstream(LocalHttpServer): super().__init__(("127.0.0.1", 0), _UpstreamHandler) self.response_status: Final = status -class _UpstreamHandler(LocalHttpHandler): +class _UpstreamHandler(LocalHttpHandler): def do_POST(self) -> None: length: Final = int(self.headers.get("content-length") or "0") self.rfile.read(length) @@ -59,6 +59,7 @@ class _UpstreamHandler(LocalHttpHandler): self.end_headers() self.wfile.write(body) + def _upstream(status: int = 200) -> AbstractContextManager[_Upstream]: return serve_in_thread(_Upstream(status)) diff --git a/tests/rust-python-harness/shared/parity/fixtures/test_recording.py b/tests/rust-python-harness/shared/parity/fixtures/test_recording.py index 6181f18e89a..c1e7c2af8d8 100644 --- a/tests/rust-python-harness/shared/parity/fixtures/test_recording.py +++ b/tests/rust-python-harness/shared/parity/fixtures/test_recording.py @@ -238,6 +238,7 @@ class _ControlledUpstreamHandler(LocalHttpHandler): self.end_headers() self.wfile.write(body) + def _controlled_upstream( stream_chunks: tuple[bytes, ...] = _SSE_CHUNKS, ) -> AbstractContextManager[_ControlledUpstream]: diff --git a/tests/rust-python-harness/shared/reporting/models.py b/tests/rust-python-harness/shared/reporting/models.py index 1ebba6c9793..4a46a6d4cfe 100644 --- a/tests/rust-python-harness/shared/reporting/models.py +++ b/tests/rust-python-harness/shared/reporting/models.py @@ -180,15 +180,11 @@ class HarnessRun: @property def unique_checks(self) -> int: - return len( - {nodeid for result in self.results.values() for nodeid in result.collected} - ) + return len({nodeid for result in self.results.values() for nodeid in result.collected}) @property def completed_checks(self) -> int: - return len( - {nodeid for result in self.results.values() for nodeid in result.completed} - ) + return len({nodeid for result in self.results.values() for nodeid in result.completed}) @classmethod def from_cases(cls, cases: Iterable[HarnessCase]) -> HarnessRun: diff --git a/tests/rust-python-harness/shared/reporting/strategy.py b/tests/rust-python-harness/shared/reporting/strategy.py index 7e76f035e20..d8e9d9e5ba9 100644 --- a/tests/rust-python-harness/shared/reporting/strategy.py +++ b/tests/rust-python-harness/shared/reporting/strategy.py @@ -67,6 +67,13 @@ class RunnerArgumentDefinition: metavar: str = "ARG" +@dataclass(frozen=True, slots=True) +class RunnerOptionDefinition: + option: str + help: str + choices: tuple[str, ...] + + class StrategyRunner(Protocol): def __call__( self, @@ -90,3 +97,4 @@ class StrategyDefinition: render: StrategyRenderer surfaces: tuple[Surface, ...] = () runner_argument: RunnerArgumentDefinition | None = None + runner_options: tuple[RunnerOptionDefinition, ...] = () diff --git a/tests/rust-python-harness/shared/test_native_build.py b/tests/rust-python-harness/shared/test_native_build.py index f3e5aead846..982f787aefa 100644 --- a/tests/rust-python-harness/shared/test_native_build.py +++ b/tests/rust-python-harness/shared/test_native_build.py @@ -89,8 +89,8 @@ def test_ensure_trace_bridge_reports_failed_rebuild(tmp_path: Final, monkeypatch assert "boom" in message -def test_ensure_trace_bridge_flags_missing_trace_feature_without_rebuild( - tmp_path: Final, monkeypatch: pytest.MonkeyPatch +def test_ensure_trace_bridge_rebuilds_when_trace_feature_is_missing( + tmp_path: Final, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] ) -> None: native: Final = tmp_path / "_native.abi3.so" native.write_bytes(b"") @@ -107,10 +107,14 @@ def test_ensure_trace_bridge_flags_missing_trace_feature_without_rebuild( monkeypatch.setattr(native_build, "_native_module_path", lambda: native) monkeypatch.setattr(native_build, "_rebuild", fake_rebuild) - monkeypatch.setattr(native_build, "get_native_bridge", lambda: SimpleNamespace(_trace=None)) + monkeypatch.setattr( + native_build, + "get_native_bridge", + lambda: SimpleNamespace(_trace=object() if state.rebuilt else None), + ) message: Final = native_build.ensure_trace_bridge(tmp_path) - assert message is not None - assert "_trace" in message - assert state.rebuilt is False + assert message is None + assert state.rebuilt is True + assert "Rebuilding native Rust bridge" in capsys.readouterr().out diff --git a/tests/rust-python-harness/shared/tracing/steps.py b/tests/rust-python-harness/shared/tracing/steps.py index 2475f0fdcc5..152f753a536 100644 --- a/tests/rust-python-harness/shared/tracing/steps.py +++ b/tests/rust-python-harness/shared/tracing/steps.py @@ -181,12 +181,7 @@ class TraceDiff: @property def matches(self) -> bool: - return ( - not self.python_only - and not self.rust_only - and not self.missing_mappings - and self.shared_order_matches - ) + return not self.python_only and not self.rust_only and not self.missing_mappings and self.shared_order_matches def _missing_mappings( @@ -257,9 +252,7 @@ def trace_diff( rust_counts: Final = Counter(rust_spans) python_only_counts: Final = python_counts - rust_counts rust_only_counts: Final = rust_counts - python_counts - python_only: Final = tuple( - span for span, count in python_only_counts.items() for _ in range(count) - ) + python_only: Final = tuple(span for span, count in python_only_counts.items() for _ in range(count)) rust_only: Final = tuple(span for span, count in rust_only_counts.items() for _ in range(count)) first_difference: Final = _first_difference(python, rust, mappings, contract) return TraceDiff( diff --git a/tests/rust-python-harness/shared/tracing/test_steps.py b/tests/rust-python-harness/shared/tracing/test_steps.py index 2efc6a3c579..1b2f02a4ed7 100644 --- a/tests/rust-python-harness/shared/tracing/test_steps.py +++ b/tests/rust-python-harness/shared/tracing/test_steps.py @@ -146,9 +146,7 @@ def test_trace_diff_allows_reordered_concurrent_children() -> None: def test_trace_diff_prunes_declared_engine_only_nodes_but_requires_them() -> None: mappings: Final = (MAPPINGS[0], mapping(rust_span="rust_prepare")) python: Final = pipeline_projection("python", (event(0, "module.py:1 entry"),), mappings).steps - rust: Final = pipeline_projection( - "rust", (event(0, "route"), event(1, "rust_prepare", 0)), mappings - ).steps + rust: Final = pipeline_projection("rust", (event(0, "route"), event(1, "rust_prepare", 0)), mappings).steps assert trace_diff(python, rust, mappings).matches assert trace_diff(python, rust[:1], mappings).missing_mappings == ("rust_prepare",) diff --git a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/fixtures/reducto.py b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/fixtures/reducto.py index bcca8ac6d42..1b283d75ca5 100644 --- a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/fixtures/reducto.py +++ b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/fixtures/reducto.py @@ -264,9 +264,7 @@ def _formatting_strategy() -> SearchStrategy[ReductoFormatting]: ), st.sampled_from((False, True)).map(lambda value: {"add_page_markers": value}), st.sampled_from((False, True)).map(lambda value: {"merge_tables": value}), - st.sampled_from(REDUCTO_FORMATTING_INCLUDE_GROUPS) - .map(list) - .map(lambda value: {"include": value}), + st.sampled_from(REDUCTO_FORMATTING_INCLUDE_GROUPS).map(list).map(lambda value: {"include": value}), ) return values.map(ReductoFormatting.model_validate) diff --git a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py index 5c4abc78081..f14296fbf06 100644 --- a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py +++ b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py @@ -124,10 +124,7 @@ class RecordingCallback(CustomLogger): if isinstance(value, Mapping): if any(not isinstance(map_key, str) for map_key in value): raise TypeError("callback kwarg mappings must use string keys") - return { - map_key: self._normalized_kwargs(map_value, map_key) - for map_key, map_value in value.items() - } + return {map_key: self._normalized_kwargs(map_value, map_key) for map_key, map_value in value.items()} if isinstance(value, (list, tuple)): return [self._normalized_kwargs(item) for item in value] raise TypeError(f"unsupported callback kwarg type: {type(value)}") diff --git a/tests/rust-python-harness/strategies/trace_parity/AGENTS.md b/tests/rust-python-harness/strategies/trace_parity/AGENTS.md index bb7cb8c91d8..5db17974d07 100644 --- a/tests/rust-python-harness/strategies/trace_parity/AGENTS.md +++ b/tests/rust-python-harness/strategies/trace_parity/AGENTS.md @@ -1 +1 @@ -Maps Python profiler frames onto feature-gated Rust span names via an explicit per-case mapping (Rust span name is the identity) and compares steps, order, and nesting of both live traces against a replayed provider response. +Prints filtered Python profiler frames and feature-gated Rust spans from live traces against a replayed provider response. The two traces are independent and are not compared. diff --git a/tests/rust-python-harness/strategies/trace_parity/__init__.py b/tests/rust-python-harness/strategies/trace_parity/__init__.py index ec88b0169fa..710bdaa3d39 100644 --- a/tests/rust-python-harness/strategies/trace_parity/__init__.py +++ b/tests/rust-python-harness/strategies/trace_parity/__init__.py @@ -7,6 +7,7 @@ from ...shared.reporting.strategy import ( ModuleCaseSpec, NotImplementedCaseSpec, RunnerArgumentDefinition, + RunnerOptionDefinition, StrategyDefinition, ) from .reporting import render_trace_results @@ -26,13 +27,17 @@ CASES: Final[tuple[CaseDefinition, ...]] = ( ModuleCaseSpec( coverage=Coverage.PARTIAL, module="tests.rust-python-harness.strategies.trace_parity.sdk.messages.case", - note="Async only until anthropic_messages_handler supports sync calls.", + note="Success paths are async; sync tracing captures the currently unsupported behavior.", ), surface="sdk", ), CaseDefinition( "responses", - NotImplementedCaseSpec(reason="No Responses trace-parity case is registered."), + ModuleCaseSpec( + coverage=Coverage.PARTIAL, + module="tests.rust-python-harness.strategies.trace_parity.sdk.responses.case", + note="Core create paths: native, streaming, provider error, Azure override, and chat bridge.", + ), surface="sdk", ), CaseDefinition( @@ -70,13 +75,17 @@ CASES: Final[tuple[CaseDefinition, ...]] = ( ModuleCaseSpec( coverage=Coverage.PARTIAL, module="tests.rust-python-harness.strategies.trace_parity.gateway.messages.case", - note="Non-streaming success paths only.", + note="Anthropic/Azure provider routes plus a fully consumed downstream streaming path.", ), surface="gateway", ), CaseDefinition( "responses", - NotImplementedCaseSpec(reason="No gateway Responses trace-parity case is registered."), + ModuleCaseSpec( + coverage=Coverage.PARTIAL, + module="tests.rust-python-harness.strategies.trace_parity.gateway.responses.case", + note="Native OpenAI non-streaming and fully consumed downstream streaming paths.", + ), surface="gateway", ), CaseDefinition( @@ -86,7 +95,11 @@ CASES: Final[tuple[CaseDefinition, ...]] = ( ), CaseDefinition( "chat_completions", - NotImplementedCaseSpec(reason="No gateway chat trace-parity case is registered."), + ModuleCaseSpec( + coverage=Coverage.PARTIAL, + module="tests.rust-python-harness.strategies.trace_parity.gateway.chat_completions.case", + note="Anthropic non-streaming and fully consumed downstream streaming paths.", + ), surface="gateway", ), CaseDefinition( @@ -99,8 +112,8 @@ CASES: Final[tuple[CaseDefinition, ...]] = ( STRATEGY: Final = StrategyDefinition( id="trace_parity", order=20, - label="Trace parity", - description="Compare pipeline steps, order, and nesting between Python profiler frames and Rust spans via an explicit mapping.", + label="Traces", + description="Print Python profiler frames and Rust spans for representative pipeline scenarios.", directory=Path(__file__).parent, runnable_spec=ModuleCaseSpec, cases=CASES, @@ -112,4 +125,11 @@ STRATEGY: Final = StrategyDefinition( metavar="NAME", help="run only this named trace scenario; repeat to select more than one", ), + runner_options=( + RunnerOptionDefinition( + option="--engine", + choices=("python", "rust"), + help="show only this engine's trace; omit to print both engines", + ), + ), ) diff --git a/tests/rust-python-harness/strategies/trace_parity/fixtures.py b/tests/rust-python-harness/strategies/trace_parity/fixtures.py new file mode 100644 index 00000000000..84ac525efb6 --- /dev/null +++ b/tests/rust-python-harness/strategies/trace_parity/fixtures.py @@ -0,0 +1,176 @@ +from __future__ import annotations + +import base64 +import binascii +import json +import struct +from collections.abc import Iterable, Mapping +from typing import Final + +from ...shared.parity.recorded_http import ( + HttpHeader, + RecordedHttpResponse, + RecordedHttpStreamResponse, + RecordedStreamChunk, +) + +JSON_HEADERS: Final = (HttpHeader(name="content-type", value="application/json"),) +SSE_HEADERS: Final = (HttpHeader(name="content-type", value="text/event-stream"),) +AWS_EVENT_STREAM_HEADERS: Final = (HttpHeader(name="content-type", value="application/vnd.amazon.eventstream"),) + + +def json_response(body: Mapping[str, object] | bytes, *, status: int = 200) -> RecordedHttpResponse: + encoded: Final = body if isinstance(body, bytes) else json.dumps(body).encode() + return RecordedHttpResponse.from_bytes(status, JSON_HEADERS, encoded) + + +def sse_event(event: str, payload: Mapping[str, object]) -> bytes: + return f"event: {event}\ndata: {json.dumps(payload, separators=(',', ':'))}\n\n".encode() + + +def sse_response(events: Iterable[tuple[str, Mapping[str, object]]]) -> RecordedHttpStreamResponse: + return RecordedHttpStreamResponse( + kind="http_stream", + status_code=200, + headers=SSE_HEADERS, + chunks=tuple(RecordedStreamChunk.from_bytes(sse_event(event, payload)) for event, payload in events), + ) + + +def _aws_string_header(name: str, value: str) -> bytes: + name_bytes: Final = name.encode() + value_bytes: Final = value.encode() + return ( + struct.pack("!B", len(name_bytes)) + + name_bytes + + struct.pack("!B", 7) + + struct.pack("!H", len(value_bytes)) + + value_bytes + ) + + +def aws_event_stream_frame(payload: Mapping[str, object]) -> bytes: + event_payload: Final = json.dumps( + {"bytes": base64.b64encode(json.dumps(payload, separators=(",", ":")).encode()).decode()}, + separators=(",", ":"), + ).encode() + headers: Final = ( + _aws_string_header(":event-type", "chunk") + + _aws_string_header(":content-type", "application/json") + + _aws_string_header(":message-type", "event") + ) + total_length: Final = 12 + len(headers) + len(event_payload) + 4 + prelude: Final = struct.pack("!II", total_length, len(headers)) + prelude_crc: Final = binascii.crc32(prelude) & 0xFFFFFFFF + prelude_crc_bytes: Final = struct.pack("!I", prelude_crc) + message_crc: Final = binascii.crc32(prelude_crc_bytes + headers + event_payload, prelude_crc) & 0xFFFFFFFF + return prelude + prelude_crc_bytes + headers + event_payload + struct.pack("!I", message_crc) + + +def aws_event_stream_response( + events: Iterable[Mapping[str, object]], *, corrupt_last_frame: bool = False +) -> RecordedHttpStreamResponse: + frames: Final = [aws_event_stream_frame(event) for event in events] + if corrupt_last_frame: + corrupted: Final = bytearray(frames[-1]) + corrupted[-1] ^= 0xFF + frames[-1] = bytes(corrupted) + return RecordedHttpStreamResponse( + kind="http_stream", + status_code=200, + headers=AWS_EVENT_STREAM_HEADERS, + chunks=(RecordedStreamChunk.from_bytes(b"".join(frames)),), + ) + + +def anthropic_response_body(*, model: str = "claude-sonnet-5") -> dict[str, object]: + return { + "id": "msg_trace", + "type": "message", + "role": "assistant", + "model": model, + "content": [{"type": "text", "text": "hello"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 2, "output_tokens": 3}, + } + + +def anthropic_stream_events(*, model: str = "claude-sonnet-5") -> tuple[tuple[str, Mapping[str, object]], ...]: + return ( + ( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_trace", + "type": "message", + "role": "assistant", + "model": model, + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 2, "output_tokens": 0}, + }, + }, + ), + ( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + ( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hello"}}, + ), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 1}, + }, + ), + ("message_stop", {"type": "message_stop"}), + ) + + +def responses_body(*, model: str = "gpt-5", status: str = "completed") -> dict[str, object]: + return { + "id": "resp_trace", + "object": "response", + "created_at": 1_750_000_000, + "status": status, + "model": model, + "output": [ + { + "type": "message", + "id": "msg_trace", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "hello", "annotations": []}], + } + ], + "usage": {"input_tokens": 2, "output_tokens": 3, "total_tokens": 5}, + } + + +def responses_stream_events(*, model: str = "gpt-5") -> tuple[tuple[str, Mapping[str, object]], ...]: + response: Final = responses_body(model=model) + return ( + ( + "response.created", + {"type": "response.created", "response": {**response, "status": "in_progress", "output": []}}, + ), + ( + "response.output_text.delta", + { + "type": "response.output_text.delta", + "item_id": "msg_trace", + "output_index": 0, + "content_index": 0, + "delta": "hello", + }, + ), + ("response.completed", {"type": "response.completed", "response": response}), + ) diff --git a/tests/rust-python-harness/strategies/trace_parity/gateway/chat_completions/case.py b/tests/rust-python-harness/strategies/trace_parity/gateway/chat_completions/case.py new file mode 100644 index 00000000000..522d2d0259c --- /dev/null +++ b/tests/rust-python-harness/strategies/trace_parity/gateway/chat_completions/case.py @@ -0,0 +1,65 @@ +from __future__ import annotations + +from typing import Final + +from .....shared.tracing.steps import Engine, mapping +from ...fixtures import anthropic_response_body, anthropic_stream_events, json_response, sse_response +from ...models import GatewayRouteSpec, RouteFixture, TraceScenario, TraceSuite + +MAPPINGS: Final = ( + mapping(span="python_chat_gateway_route", python_frame=r"proxy_server\.py:\d+ chat_completion$"), + mapping(span="python_gateway_service", python_frame=r"ProxyBaseLLMRequestProcessing\.base_process_llm_request$"), + mapping(rust_span="chat_completions_gateway_route"), + mapping(rust_span="chat_completions"), + mapping(span="python_chat_entrypoint", python_frame=r"main\.py:\d+ a?completion$"), + mapping(span="python_provider_config", python_frame=r"ProviderConfigManager\.get_provider_chat_config$"), + mapping(rust_span="validate_environment", python_frame=r"(? RouteFixture: + return RouteFixture( + kwargs={ + "model_alias": "trace-model", + "provider_model": "anthropic/claude-sonnet-5", + "body": { + "model": "trace-model", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 16, + }, + }, + provider_responses=(json_response(anthropic_response_body()),), + ) + + +def _stream_fixture(engine: Engine, base_url: str) -> RouteFixture: + fixture: Final = _fixture(engine, base_url) + return fixture.with_body(stream=True).derive( + provider_responses=(sse_response(anthropic_stream_events()),), + ) + + +TRACE_SUITE: Final = TraceSuite( + route=GatewayRouteSpec("chat_completions"), + scenarios=( + TraceScenario(name="async-anthropic", fixture=_fixture, mappings=MAPPINGS, asynchronous=True), + TraceScenario( + name="async-anthropic-downstream-stream", + fixture=_stream_fixture, + mappings=(*MAPPINGS, *STREAM_MAPPINGS), + asynchronous=True, + ), + ), +) diff --git a/tests/rust-python-harness/strategies/trace_parity/gateway/execution.py b/tests/rust-python-harness/strategies/trace_parity/gateway/execution.py index 860e872dd44..88130dab808 100644 --- a/tests/rust-python-harness/strategies/trace_parity/gateway/execution.py +++ b/tests/rust-python-harness/strategies/trace_parity/gateway/execution.py @@ -13,8 +13,8 @@ from ....shared.parity.replay import replay_server from ....shared.tracing.native import TraceResponsePayload, native_trace_events from ....shared.tracing.profiler import FunctionTraceEvent, profile_python from ....shared.tracing.steps import Engine, PipelineProjection, pipeline_projection -from ..models import GatewayRouteSpec, RouteFixture, TraceExecutionFailure, TraceMode, TraceScenario -from ..reporting import TraceComparisonArtifact +from ..models import GatewayRouteSpec, RouteFixture, TraceEngine, TraceExecutionFailure, TraceScenario +from ..reporting import TraceArtifact class _GatewayResponsePayload(BaseModel): @@ -28,7 +28,14 @@ class _GatewayClient(Protocol): def post(self, url: str, *, json: object, headers: dict[str, str]) -> httpx.Response: ... -def _collect_python(fixture: RouteFixture) -> tuple[FunctionTraceEvent, ...]: +_ROUTE_PATHS: Final = { + "messages": "/v1/messages", + "chat_completions": "/v1/chat/completions", + "responses": "/v1/responses", +} + + +def _collect_python(fixture: RouteFixture, route: GatewayRouteSpec) -> tuple[FunctionTraceEvent, ...]: from fastapi.testclient import TestClient import litellm @@ -61,7 +68,7 @@ def _collect_python(fixture: RouteFixture) -> tuple[FunctionTraceEvent, ...]: with profile_python(Path(litellm.__file__).parent, threads=True) as profiler: client: Final = cast(_GatewayClient, TestClient(proxy_server.app)) response: Final = client.post( - "/v1/messages", + _ROUTE_PATHS[route.route], json=fixture.kwargs["body"], headers={"authorization": "Bearer trace-key"}, ) @@ -76,9 +83,10 @@ def _collect_python(fixture: RouteFixture) -> tuple[FunctionTraceEvent, ...]: proxy_server.app.dependency_overrides[user_api_key_auth] = old_override -def _collect_rust(fixture: RouteFixture) -> tuple[FunctionTraceEvent, ...]: +def _collect_rust(fixture: RouteFixture, route: GatewayRouteSpec) -> tuple[FunctionTraceEvent, ...]: payload: Final = json.dumps( { + "path": _ROUTE_PATHS[route.route], "model_alias": fixture.kwargs["model_alias"], "provider_model": fixture.kwargs["provider_model"], "api_base": fixture.kwargs["api_base"], @@ -130,7 +138,9 @@ def _gateway_trace_binary() -> Path: return rust_root / "target" / "debug" / "trace-parity-gateway" -def _collect(scenario: TraceScenario, engine: Engine) -> tuple[FunctionTraceEvent, ...] | TraceExecutionFailure: +def _collect( + route: GatewayRouteSpec, scenario: TraceScenario, engine: Engine +) -> tuple[FunctionTraceEvent, ...] | TraceExecutionFailure: try: with replay_server() as provider: base_fixture: Final = scenario.fixture(engine, provider.url) @@ -140,7 +150,7 @@ def _collect(scenario: TraceScenario, engine: Engine) -> tuple[FunctionTraceEven ) for response in fixture.provider_responses: provider.enqueue_response(response) - events: Final = _collect_python(fixture) if engine == "python" else _collect_rust(fixture) + events: Final = _collect_python(fixture, route) if engine == "python" else _collect_rust(fixture, route) provider.take_requests(len(fixture.provider_responses)) return events except Exception as error: @@ -151,9 +161,8 @@ def _projections( python_events: tuple[FunctionTraceEvent, ...], rust_events: tuple[FunctionTraceEvent, ...], scenario: TraceScenario, - mode: TraceMode, ) -> tuple[PipelineProjection, PipelineProjection, str | None]: - mappings: Final = scenario.mappings_for(mode) + mappings: Final = scenario.mappings try: return ( pipeline_projection("python", python_events, mappings), @@ -164,26 +173,26 @@ def _projections( return PipelineProjection(), PipelineProjection(), f"harness: {error}" -def execute_gateway_trace(route: GatewayRouteSpec, scenario: TraceScenario, mode: TraceMode) -> TraceComparisonArtifact: - mappings: Final = scenario.mappings_for(mode) - python_trace: Final = _collect(scenario, "python") - rust_trace: Final = _collect(scenario, "rust") +def execute_gateway_trace( + route: GatewayRouteSpec, + scenario: TraceScenario, + engine: TraceEngine = "both", +) -> TraceArtifact: + python_trace: Final = _collect(route, scenario, "python") if engine != "rust" else () + rust_trace: Final = _collect(route, scenario, "rust") if engine != "python" else () collection_python_error: Final = None if isinstance(python_trace, tuple) else f"python: {python_trace.message}" rust_error: Final = None if isinstance(rust_trace, tuple) else f"rust: {rust_trace.message}" python_events: Final = python_trace if isinstance(python_trace, tuple) else () rust_events: Final = rust_trace if isinstance(rust_trace, tuple) else () - python, rust, projection_error = _projections(python_events, rust_events, scenario, mode) + python, rust, projection_error = _projections(python_events, rust_events, scenario) python_error: Final = projection_error or collection_python_error - return TraceComparisonArtifact.from_traces( + return TraceArtifact.from_traces( + engine=engine, surface="gateway", sdk_function=route.route, scenario=scenario.name, - mode=mode, - mappings=mappings, - contract=scenario.contract, python=python.steps, rust=rust.steps, - python_unmatched=python.unmatched, python_error=python_error, rust_error=rust_error, ) diff --git a/tests/rust-python-harness/strategies/trace_parity/gateway/messages/case.py b/tests/rust-python-harness/strategies/trace_parity/gateway/messages/case.py index 30f51cee353..ca9c858f6b7 100644 --- a/tests/rust-python-harness/strategies/trace_parity/gateway/messages/case.py +++ b/tests/rust-python-harness/strategies/trace_parity/gateway/messages/case.py @@ -1,10 +1,9 @@ from __future__ import annotations -import json from typing import Final -from .....shared.parity.recorded_http import HttpHeader, RecordedHttpResponse from .....shared.tracing.steps import Engine, mapping +from ...fixtures import anthropic_response_body, anthropic_stream_events, json_response, sse_response from ...models import GatewayRouteSpec, RouteFixture, TraceScenario, TraceSuite @@ -49,24 +48,7 @@ def _fixture(_engine: Engine, provider: str) -> RouteFixture: "max_tokens": 16, }, }, - provider_responses=( - RecordedHttpResponse.from_bytes( - 200, - (HttpHeader(name="content-type", value="application/json"),), - json.dumps( - { - "id": "msg_trace", - "type": "message", - "role": "assistant", - "model": "claude-sonnet-5", - "content": [{"type": "text", "text": "hello"}], - "stop_reason": "end_turn", - "stop_sequence": None, - "usage": {"input_tokens": 2, "output_tokens": 3}, - } - ).encode(), - ), - ), + provider_responses=(json_response(anthropic_response_body()),), ) @@ -78,6 +60,13 @@ def _azure_fixture(engine: Engine, _base_url: str) -> RouteFixture: return _fixture(engine, "azure_ai") +def _stream_fixture(engine: Engine, _base_url: str) -> RouteFixture: + fixture: Final = _anthropic_fixture(engine, _base_url) + return fixture.with_body(stream=True).derive( + provider_responses=(sse_response(anthropic_stream_events()),), + ) + + ANTHROPIC_MAPPINGS: Final = ( *GATEWAY_MAPPINGS, mapping( @@ -100,7 +89,22 @@ AZURE_MAPPINGS: Final = ( TRACE_SUITE: Final = TraceSuite( route=GatewayRouteSpec("messages"), scenarios=( - TraceScenario(name="anthropic", fixture=_anthropic_fixture, mappings=ANTHROPIC_MAPPINGS, modes=("async",)), - TraceScenario(name="azure-ai", fixture=_azure_fixture, mappings=AZURE_MAPPINGS, modes=("async",)), + TraceScenario( + name="async-anthropic", fixture=_anthropic_fixture, mappings=ANTHROPIC_MAPPINGS, asynchronous=True + ), + TraceScenario(name="async-azure-ai", fixture=_azure_fixture, mappings=AZURE_MAPPINGS, asynchronous=True), + TraceScenario( + name="async-anthropic-downstream-stream", + fixture=_stream_fixture, + mappings=( + *ANTHROPIC_MAPPINGS, + mapping(span="python_upstream_stream", python_frame=r"AnthropicMessagesStreamingResponse\.__anext__$"), + mapping( + span="python_downstream_stream", python_frame=r"DataGenerator\.__anext__$|async_data_generator$" + ), + mapping(span="python_stream_callback", python_frame=r"Logging\.async_success_handler$"), + ), + asynchronous=True, + ), ), ) diff --git a/tests/rust-python-harness/strategies/trace_parity/gateway/responses/case.py b/tests/rust-python-harness/strategies/trace_parity/gateway/responses/case.py new file mode 100644 index 00000000000..52c1ebb391f --- /dev/null +++ b/tests/rust-python-harness/strategies/trace_parity/gateway/responses/case.py @@ -0,0 +1,63 @@ +from __future__ import annotations + +from typing import Final + +from .....shared.tracing.steps import Engine, mapping +from ...fixtures import json_response, responses_body, responses_stream_events, sse_response +from ...models import GatewayRouteSpec, RouteFixture, TraceScenario, TraceSuite + +MAPPINGS: Final = ( + mapping( + span="python_responses_gateway_route", python_frame=r"response_api_endpoints/endpoints\.py:\d+ responses_api$" + ), + mapping(span="python_gateway_service", python_frame=r"ProxyBaseLLMRequestProcessing\.base_process_llm_request$"), + mapping(rust_span="responses_gateway_route"), + mapping(span="python_responses", python_frame=r"responses/main\.py:\d+ a?responses$"), + mapping(span="python_provider_config", python_frame=r"ProviderConfigManager\.get_provider_responses_api_config$"), + mapping(rust_span="validate_environment", python_frame=r"OpenAIResponsesAPIConfig\.validate_environment$"), + mapping(rust_span="complete_url", python_frame=r"OpenAIResponsesAPIConfig\.get_complete_url$"), + mapping(rust_span="transform_request", python_frame=r"OpenAIResponsesAPIConfig\.transform_responses_api_request$"), + mapping(span="python_logging_pre_call", python_frame=r"Logging\.pre_call$"), + mapping(rust_span="http_request", python_frame=r"AsyncHTTPHandler\.post$|HTTPHandler\.post$"), + mapping(rust_span="transform_response", python_frame=r"OpenAIResponsesAPIConfig\.transform_response_api_response$"), + mapping(span="python_success_callback", python_frame=r"Logging\.async_success_handler$|Logging\.success_handler$"), +) + +STREAM_MAPPINGS: Final = ( + mapping(span="python_stream_iterator", python_frame=r"ResponsesAPIStreamingIterator\.__init__$"), + mapping(span="python_stream_next", python_frame=r"ResponsesAPIStreamingIterator\.__anext__$"), + mapping(span="python_stream_transform", python_frame=r"OpenAIResponsesAPIConfig\.transform_streaming_response$"), + mapping(span="python_downstream_stream", python_frame=r"DataGenerator\.__anext__$|async_data_generator$"), +) + + +def _fixture(_engine: Engine, _base_url: str) -> RouteFixture: + return RouteFixture( + kwargs={ + "model_alias": "trace-model", + "provider_model": "openai/gpt-5", + "body": {"model": "trace-model", "input": "hello"}, + }, + provider_responses=(json_response(responses_body()),), + ) + + +def _stream_fixture(engine: Engine, base_url: str) -> RouteFixture: + fixture: Final = _fixture(engine, base_url) + return fixture.with_body(stream=True).derive( + provider_responses=(sse_response(responses_stream_events()),), + ) + + +TRACE_SUITE: Final = TraceSuite( + route=GatewayRouteSpec("responses"), + scenarios=( + TraceScenario(name="async-openai", fixture=_fixture, mappings=MAPPINGS, asynchronous=True), + TraceScenario( + name="async-openai-downstream-stream", + fixture=_stream_fixture, + mappings=(*MAPPINGS, *STREAM_MAPPINGS), + asynchronous=True, + ), + ), +) diff --git a/tests/rust-python-harness/strategies/trace_parity/models.py b/tests/rust-python-harness/strategies/trace_parity/models.py index 04659b25382..078af2b316d 100644 --- a/tests/rust-python-harness/strategies/trace_parity/models.py +++ b/tests/rust-python-harness/strategies/trace_parity/models.py @@ -1,22 +1,45 @@ from __future__ import annotations -from collections.abc import Callable +from collections.abc import Callable, Mapping from dataclasses import dataclass -from typing import Final, Literal, TypeAlias +from typing import Final, Literal, TypeAlias, cast -from ...shared.parity.recorded_http import RecordedHttpResponse +from ...shared.parity.recorded_http import RecordedResponse from ...shared.reporting.models import SdkFunction -from ...shared.tracing.steps import Engine, TraceContract, TraceMapping +from ...shared.tracing.steps import Engine, TraceMapping -TraceMode = Literal["sync", "async"] +TraceEngine = Literal["python", "rust", "both"] TraceFailureSource = Literal["python", "rust", "harness"] @dataclass(frozen=True, slots=True) class RouteFixture: kwargs: dict[str, object] - provider_responses: tuple[RecordedHttpResponse, ...] + provider_responses: tuple[RecordedResponse, ...] expected_failure: bool = False + consume_stream: bool = False + + def derive( + self, + *, + kwargs: Mapping[str, object] | None = None, + provider_responses: tuple[RecordedResponse, ...] | None = None, + expected_failure: bool | None = None, + consume_stream: bool | None = None, + ) -> RouteFixture: + return RouteFixture( + kwargs={**self.kwargs, **(kwargs or {})}, + provider_responses=self.provider_responses if provider_responses is None else provider_responses, + expected_failure=self.expected_failure if expected_failure is None else expected_failure, + consume_stream=self.consume_stream if consume_stream is None else consume_stream, + ) + + def with_body(self, **updates: object) -> RouteFixture: + raw_body: Final = self.kwargs.get("body") + if not isinstance(raw_body, dict): + raise ValueError("route fixture does not contain an object body") + body: Final = cast(dict[str, object], raw_body) + return self.derive(kwargs={"body": {**body, **updates}}) @dataclass(frozen=True, slots=True) @@ -40,14 +63,7 @@ class TraceScenario: name: str fixture: Callable[[Engine, str], RouteFixture] mappings: tuple[TraceMapping, ...] - modes: tuple[TraceMode, ...] = ("sync", "async") - contract: TraceContract = TraceContract() - sync_mappings: tuple[TraceMapping, ...] | None = None - async_mappings: tuple[TraceMapping, ...] | None = None - - def mappings_for(self, mode: TraceMode) -> tuple[TraceMapping, ...]: - selected: Final = self.async_mappings if mode == "async" else self.sync_mappings - return self.mappings if selected is None else selected + asynchronous: bool @dataclass(frozen=True, slots=True) diff --git a/tests/rust-python-harness/strategies/trace_parity/reporting.py b/tests/rust-python-harness/strategies/trace_parity/reporting.py index 9c5bf9e88cd..e7c07ef9c0f 100644 --- a/tests/rust-python-harness/strategies/trace_parity/reporting.py +++ b/tests/rust-python-harness/strategies/trace_parity/reporting.py @@ -1,31 +1,24 @@ from __future__ import annotations import os -import re import sys from collections.abc import Sequence -from typing import Final, Literal +from typing import Final from pydantic import BaseModel, ConfigDict, ValidationError from ...shared.reporting.models import SURFACES, CaseResult, RunStatus, SdkFunction, Surface from ...shared.reporting.rendering import ReportSection from ...shared.reporting.strategy import NotImplementedCaseSpec, SkippedCaseSpec -from ...shared.tracing.steps import ( - PipelineStep, - TraceContract, - TraceDiff, - TraceMapping, - trace_depths, - trace_diff, -) +from ...shared.tracing.steps import PipelineStep, trace_depths +from .models import TraceEngine -TRACE_COMPARISON_ARTIFACT: Final = "trace_comparison" +TRACE_ARTIFACT: Final = "trace" TRACE_PARITY_HINT: Final = ( "rebuild the native bridge with the trace-parity feature, e.g. `uvx maturin develop --features trace-parity`" ) -_COLORS: Final[dict[str, str]] = {"green": "32", "yellow": "33", "red": "31", "cyan": "36"} +_COLORS: Final[dict[str, str]] = {"yellow": "33", "red": "31", "cyan": "36"} _RESET: Final = "\033[0m" @@ -47,26 +40,15 @@ class TraceEventArtifact(BaseModel): return PipelineStep(self.id, self.parent_id, self.span, self.raw) -class TraceMappingArtifact(BaseModel): - model_config = ConfigDict(frozen=True, extra="forbid") - - span: str - python: str | None - rust: str | None - - -class TraceComparisonArtifact(BaseModel): +class TraceArtifact(BaseModel): model_config = ConfigDict(frozen=True, extra="forbid") + engine: TraceEngine = "both" surface: Surface sdk_function: SdkFunction scenario: str - mode: Literal["sync", "async"] - mappings: tuple[TraceMappingArtifact, ...] python: tuple[TraceEventArtifact, ...] rust: tuple[TraceEventArtifact, ...] - python_unmatched: int - unordered_children_of: frozenset[str] python_error: str | None = None rust_error: str | None = None @@ -74,41 +56,27 @@ class TraceComparisonArtifact(BaseModel): def from_traces( cls, *, + engine: TraceEngine = "both", surface: Surface, sdk_function: SdkFunction, scenario: str, - mode: Literal["sync", "async"], - mappings: Sequence[TraceMapping], - contract: TraceContract, python: Sequence[PipelineStep], rust: Sequence[PipelineStep], - python_unmatched: int, python_error: str | None = None, rust_error: str | None = None, - ) -> TraceComparisonArtifact: + ) -> TraceArtifact: return cls( + engine=engine, surface=surface, sdk_function=sdk_function, scenario=scenario, - mode=mode, - mappings=tuple( - TraceMappingArtifact( - span=item.span, - python=item.python.pattern if item.python else None, - rust=item.rust, - ) - for item in mappings - ), python=tuple( TraceEventArtifact(id=step.id, parent_id=step.parent_id, span=step.span, raw=step.raw) for step in python ), rust=tuple( - TraceEventArtifact(id=step.id, parent_id=step.parent_id, span=step.span, raw=step.raw) - for step in rust + TraceEventArtifact(id=step.id, parent_id=step.parent_id, span=step.span, raw=step.raw) for step in rust ), - python_unmatched=python_unmatched, - unordered_children_of=contract.unordered_children_of, python_error=python_error, rust_error=rust_error, ) @@ -119,32 +87,9 @@ class TraceComparisonArtifact(BaseModel): def rust_steps(self) -> tuple[PipelineStep, ...]: return tuple(event.step() for event in self.rust) - def diff(self) -> TraceDiff: - return trace_diff( - self.python_steps(), - self.rust_steps(), - tuple( - TraceMapping( - item.span, - re.compile(item.python) if item.python is not None else None, - item.rust, - ) - for item in self.mappings - ), - TraceContract(self.unordered_children_of), - ) - - def exact_match(self) -> bool: - return self.diff().matches - def has_errors(self) -> bool: return self.python_error is not None or self.rust_error is not None - def contract_matches(self) -> bool: - if self.has_errors(): - return False - return self.diff().matches - def _split_raw(raw: str) -> tuple[str, str]: location, separator, name = raw.partition(" ") @@ -153,72 +98,28 @@ def _split_raw(raw: str) -> tuple[str, str]: return raw, "" -def _python_line(index: int, step: PipelineStep, depth: int, exclusive: frozenset[str]) -> str: +def _python_line(index: int, step: PipelineStep, depth: int) -> str: name: Final = _split_raw(step.raw)[0] location: Final = _split_raw(step.raw)[1] suffix: Final = f" ({location})" if location else "" - marker: Final = " [python only]" if step.span in exclusive else "" - return _paint(f"{index} {' ' * depth}{name}{suffix}{marker}", "cyan") + return _paint(f"{index} {' ' * depth}{name}{suffix}", "cyan") -def _python_lines(steps: tuple[PipelineStep, ...], exclusive: frozenset[str]) -> str: +def _python_lines(steps: tuple[PipelineStep, ...]) -> str: depths: Final = trace_depths(steps) - lines: Final = tuple( - _python_line(index, step, depths[step.id], exclusive) for index, step in enumerate(steps, start=1) - ) + lines: Final = tuple(_python_line(index, step, depths[step.id]) for index, step in enumerate(steps, start=1)) return f"{_paint('PYTHON', 'cyan')} ({len(steps)} steps)\n" + ("\n".join(lines) if lines else "(empty)") -def _python_references(steps: tuple[PipelineStep, ...]) -> dict[tuple[str, int], str]: - references: dict[tuple[str, int], str] = {} - occurrences: dict[str, int] = {} - for index, step in enumerate(steps, start=1): - name = _split_raw(step.raw)[0] - occurrence = occurrences.get(step.span, 0) + 1 - occurrences[step.span] = occurrence - references[(step.span, occurrence)] = f"{index} {name}" - return references - - -def _rust_line( - step: PipelineStep, - depth: int, - occurrence: int, - references: dict[tuple[str, int], str], -) -> str: - span: Final = _paint(step.span, "yellow") - key: Final = (step.span, occurrence) - reference: Final = ( - _paint(references[key], "cyan") if key in references else _paint("[rust only]", "yellow") - ) - suffix: Final = f"#{occurrence}" if occurrence > 1 else "" - return f"{' ' * depth}{span}{suffix} -> {reference}" - - -def _rust_lines(steps: tuple[PipelineStep, ...], references: dict[tuple[str, int], str]) -> str: +def _rust_lines(steps: tuple[PipelineStep, ...]) -> str: depths: Final = trace_depths(steps) - occurrences: dict[str, int] = {} - lines: list[str] = [] - for step in steps: - occurrence = occurrences.get(step.span, 0) + 1 - occurrences[step.span] = occurrence - lines.append(_rust_line(step, depths[step.id], occurrence, references)) + lines: Final = tuple( + _paint(f"{index} {' ' * depths[step.id]}{step.span}", "yellow") for index, step in enumerate(steps, 1) + ) return f"{_paint('RUST', 'yellow')} ({len(steps)} steps)\n" + ("\n".join(lines) if lines else "(empty)") -def _state_text(state: str, *, good: bool) -> str: - return _paint(state, "green" if good else "red") - - -def _contract_line(artifact: TraceComparisonArtifact) -> str: - matches: Final = artifact.contract_matches() - status: Final = _state_text("PASS" if matches else "FAIL", good=matches) - if artifact.python_error or artifact.rust_error: - return f"Contract: {status}" - return f"Contract: {status}" - - -def _error_lines(artifact: TraceComparisonArtifact) -> tuple[str, ...]: +def _error_lines(artifact: TraceArtifact) -> tuple[str, ...]: lines: list[str] = [] for engine, error in (("Python", artifact.python_error), ("Rust", artifact.rust_error)): if error is None: @@ -229,68 +130,20 @@ def _error_lines(artifact: TraceComparisonArtifact) -> tuple[str, ...]: return tuple(lines) -def _unseen_mappings( - artifact: TraceComparisonArtifact, - python: tuple[PipelineStep, ...], - rust: tuple[PipelineStep, ...], -) -> tuple[str, ...]: - return artifact.diff().missing_mappings - - -def _comparison_status_lines( - artifact: TraceComparisonArtifact, - python: tuple[PipelineStep, ...], - rust: tuple[PipelineStep, ...], -) -> tuple[str, ...]: - diff: Final = artifact.diff() - exact_match: Final = artifact.exact_match() - if artifact.has_errors(): - return (*_error_lines(artifact), _contract_line(artifact)) - unseen: Final = _unseen_mappings(artifact, python, rust) - unseen_line: Final[tuple[str, ...]] = (f"Unseen mappings: {', '.join(unseen)}",) if unseen else () - drift_lines: Final[tuple[str, ...]] = ( - (_state_text("Same steps, order, and nesting", good=True),) - if exact_match - else ( - _paint(f"Python only: {', '.join(diff.python_only) or 'none'}", "cyan"), - _paint(f"Rust only: {', '.join(diff.rust_only) or 'none'}", "yellow"), - f"First difference: {diff.first_difference or 'none'}", - f"Python frames outside mapping: {artifact.python_unmatched}", - ) - ) - return ( - f"Trace: {_state_text('MATCH' if exact_match else 'DRIFT', good=exact_match)}", - *drift_lines, - *unseen_line, - _contract_line(artifact), - ) - - -def _render_comparison(artifact: TraceComparisonArtifact) -> str: - python: Final = artifact.python_steps() - rust: Final = artifact.rust_steps() - diff: Final = artifact.diff() - python_exclusive: Final = frozenset(item.span for item in artifact.mappings if item.rust is None) - status_lines: Final = _comparison_status_lines(artifact, python, rust) - return "\n\n".join( - ( - _python_lines(python, python_exclusive | frozenset(diff.python_only)), - _rust_lines(rust, _python_references(python)), - "\n".join(status_lines), - ) - ) - - -def _mode(nodeid: str) -> str: - if "[" in nodeid: - return nodeid.rsplit("[", 1)[-1].removesuffix("]") - head, _, tail = nodeid.rpartition(":") - return tail if head else "unknown mode" +def _render_trace(artifact: TraceArtifact) -> str: + traces: tuple[str, ...] + if artifact.engine == "python": + traces = (_python_lines(artifact.python_steps()),) + elif artifact.engine == "rust": + traces = (_rust_lines(artifact.rust_steps()),) + else: + traces = (_python_lines(artifact.python_steps()), _rust_lines(artifact.rust_steps())) + return "\n\n".join((*traces, *_error_lines(artifact))) def _scenario(nodeid: str) -> str: parts: Final = nodeid.split(":") - return parts[-2] if len(parts) >= 5 else "default" + return parts[-1] if len(parts) >= 4 else "default" def _unavailable(status: RunStatus) -> str: @@ -299,20 +152,20 @@ def _unavailable(status: RunStatus) -> str: def _render_artifact(body: str) -> str: try: - artifact: Final = TraceComparisonArtifact.model_validate_json(body) + artifact: Final = TraceArtifact.model_validate_json(body) except ValidationError as error: - return f"Trace comparison artifact is invalid: {error}" - return _render_comparison(artifact) + return f"Trace artifact is invalid: {error}" + return _render_trace(artifact) -def _mode_section(result: CaseResult, nodeid: str, status: RunStatus) -> str: +def _scenario_section(result: CaseResult, nodeid: str, status: RunStatus) -> str: artifacts: Final = tuple( - artifact for artifact in result.artifacts.get(nodeid, ()) if artifact.kind == TRACE_COMPARISON_ARTIFACT + artifact for artifact in result.artifacts.get(nodeid, ()) if artifact.kind == TRACE_ARTIFACT ) body: Final = ( "\n\n".join(_render_artifact(artifact.body) for artifact in artifacts) if artifacts else _unavailable(status) ) - label: Final = f"Scenario: {_scenario(nodeid)} / Mode: {_mode(nodeid)}" + label: Final = f"Scenario: {_scenario(nodeid)}" return f"{label}\n{'-' * len(label)}\n\n{body}" @@ -321,7 +174,7 @@ def _case_block(result: CaseResult) -> str: outcomes: Final = tuple(result.outcomes.items()) or ( (nodeid, RunStatus.NOT_RUN) for nodeid in sorted(result.collected) ) - sections: Final = tuple(_mode_section(result, nodeid, status) for nodeid, status in outcomes) + sections: Final = tuple(_scenario_section(result, nodeid, status) for nodeid, status in outcomes) return "\n\n".join((f"{header}\n{'=' * len(header)}", *sections)) @@ -357,11 +210,11 @@ def _surface_section(surface: Surface, results: Sequence[CaseResult]) -> ReportS *((not_implemented,) if not_implemented else ()), *((skipped,) if skipped else ()), ) - return ReportSection(f"{surface.upper()} trace comparisons", blocks or ("No runnable trace comparisons",)) + return ReportSection(f"{surface.upper()} traces", blocks or ("No runnable traces",)) def render_trace_results(results: Sequence[CaseResult]) -> tuple[ReportSection, ...]: sections: Final = tuple( section for surface in SURFACES if (section := _surface_section(surface, results)) is not None ) - return sections or (ReportSection("Trace comparisons", ("No trace comparisons selected",)),) + return sections or (ReportSection("Traces", ("No traces selected",)),) diff --git a/tests/rust-python-harness/strategies/trace_parity/runner.py b/tests/rust-python-harness/strategies/trace_parity/runner.py index b78a3c7da3f..706e054bf52 100644 --- a/tests/rust-python-harness/strategies/trace_parity/runner.py +++ b/tests/rust-python-harness/strategies/trace_parity/runner.py @@ -4,13 +4,20 @@ import importlib from collections.abc import Sequence from pathlib import Path from time import monotonic -from typing import Final +from typing import Final, cast +from ...shared.native_build import ensure_trace_bridge from ...shared.reporting.models import CaseResult, HarnessCase, HarnessRun, ResultArtifact, RunStatus, Surface from ...shared.reporting.strategy import ModuleCaseSpec, UpdateCallback -from ...shared.native_build import ensure_trace_bridge -from .models import GatewayRouteSpec, RouteSpec, TraceExecutionFailure, TraceMode, TraceScenario, TraceSuite -from .reporting import TRACE_COMPARISON_ARTIFACT, TraceComparisonArtifact +from .models import ( + GatewayRouteSpec, + RouteSpec, + TraceEngine, + TraceExecutionFailure, + TraceScenario, + TraceSuite, +) +from .reporting import TRACE_ARTIFACT, TraceArtifact from .sdk.execution import execute_trace @@ -32,15 +39,13 @@ def validate_trace_suite(suite: TraceSuite, harness_case: HarnessCase) -> str | names: Final = tuple(scenario.name for scenario in suite.scenarios) if not names or len(names) != len(set(names)) or any(not name or ":" in name for name in names): return "scenario names must be non-empty, unique, and colon-free" - invalid_modes: Final = tuple( + invalid_names: Final = tuple( scenario.name for scenario in suite.scenarios - if not scenario.modes - or len(scenario.modes) != len(set(scenario.modes)) - or any(mode not in {"sync", "async"} for mode in scenario.modes) + if not scenario.name.startswith("async-" if scenario.asynchronous else "sync-") ) - if invalid_modes: - return f"scenarios must use non-empty, unique sync/async modes: {', '.join(invalid_modes)}" + if invalid_names: + return f"scenario names must start with sync- or async-: {', '.join(invalid_names)}" surface: Final = harness_case.surface if surface == "sdk" and not isinstance(suite.route, RouteSpec): return "must use RouteSpec for the sdk surface" @@ -57,15 +62,14 @@ def scenario_nodeids( trace_suite: TraceSuite, harness_case: HarnessCase, selected_scenarios: frozenset[str] = frozenset(), -) -> tuple[tuple[TraceScenario, TraceMode, str], ...]: +) -> tuple[tuple[TraceScenario, str], ...]: surface: Final = harness_case.surface if surface is None: return () return tuple( - (scenario, mode, f"trace:{surface}:{harness_case.sdk_function}:{scenario.name}:{mode}") + (scenario, f"trace:{surface}:{harness_case.sdk_function}:{scenario.name}") for scenario in trace_suite.scenarios if not selected_scenarios or scenario.name in selected_scenarios - for mode in scenario.modes ) @@ -77,54 +81,49 @@ def _record_setup_failure(run: HarnessRun, case: HarnessCase, message: str, stag run.failures.append((nodeid, message)) -def run_trace_mode( +def run_trace_scenario( run: HarnessRun, result: CaseResult, trace_suite: TraceSuite, scenario: TraceScenario, - mode: TraceMode, surface: Surface, nodeid: str, on_update: UpdateCallback, + engine: TraceEngine = "both", ) -> None: started_at: Final = monotonic() - comparison: Final = _execute_mode(trace_suite, scenario, mode, surface) + trace: Final = _execute_scenario(trace_suite, scenario, surface, engine) duration: Final = monotonic() - started_at - if isinstance(comparison, TraceExecutionFailure): + if isinstance(trace, TraceExecutionFailure): result.record(nodeid, RunStatus.ERROR, duration) - run.failures.append((nodeid, comparison.message)) + run.failures.append((nodeid, trace.message)) on_update(run) return - artifact: Final = ResultArtifact(TRACE_COMPARISON_ARTIFACT, comparison.model_dump_json()) - if comparison.has_errors(): + artifact: Final = ResultArtifact(TRACE_ARTIFACT, trace.model_dump_json()) + if trace.has_errors(): result.record(nodeid, RunStatus.ERROR, duration, (artifact,)) - run.failures.append( - (nodeid, "\n".join(error for error in (comparison.python_error, comparison.rust_error) if error)) - ) + run.failures.append((nodeid, "\n".join(error for error in (trace.python_error, trace.rust_error) if error))) else: - status: Final = RunStatus.PASSED if comparison.contract_matches() else RunStatus.FAILED - result.record(nodeid, status, duration, (artifact,)) - if status is RunStatus.FAILED: - run.failures.append((nodeid, "trace contract mismatch; see the rendered comparison")) + result.record(nodeid, RunStatus.PASSED, duration, (artifact,)) on_update(run) -def _execute_mode( +def _execute_scenario( trace_suite: TraceSuite, scenario: TraceScenario, - mode: TraceMode, surface: Surface, -) -> TraceComparisonArtifact | TraceExecutionFailure: + engine: TraceEngine, +) -> TraceArtifact | TraceExecutionFailure: route: Final = trace_suite.route if isinstance(route, GatewayRouteSpec): if surface != "gateway": return TraceExecutionFailure("harness", "gateway route cannot run on the sdk surface") from .gateway.execution import execute_gateway_trace - return execute_gateway_trace(route, scenario, mode) + return execute_gateway_trace(route, scenario, engine) if surface != "sdk": return TraceExecutionFailure("harness", "sdk route cannot run on the gateway surface") - return execute_trace(route, scenario, mode, surface) + return execute_trace(route, scenario, surface, engine) def _run_case( @@ -132,6 +131,7 @@ def _run_case( harness_case: HarnessCase, selected_scenarios: frozenset[str], on_update: UpdateCallback, + engine: TraceEngine, ) -> None: result: Final = run.results[harness_case.key] spec: Final = harness_case.spec @@ -146,15 +146,29 @@ def _run_case( on_update(run) return nodeids: Final = scenario_nodeids(trace_suite, harness_case, selected_scenarios) - result.collected.update(nodeid for _, _, nodeid in nodeids) + result.collected.update(nodeid for _, nodeid in nodeids) if not nodeids: result.status = RunStatus.SKIPPED on_update(run) return result.status = RunStatus.RUNNING on_update(run) - for scenario, mode, nodeid in nodeids: - run_trace_mode(run, result, trace_suite, scenario, mode, surface, nodeid, on_update) + for scenario, nodeid in nodeids: + run_trace_scenario(run, result, trace_suite, scenario, surface, nodeid, on_update, engine) + + +def runner_selection(runner_args: Sequence[str]) -> tuple[frozenset[str], TraceEngine]: + engine: TraceEngine = "both" + scenarios: list[str] = [] + for argument in runner_args: + if argument.startswith("--engine="): + value = argument.removeprefix("--engine=") + if value not in {"python", "rust"}: + raise ValueError(f"invalid trace engine: {value}") + engine = cast(TraceEngine, value) + else: + scenarios.append(argument) + return frozenset(scenarios), engine def run_trace_cases( @@ -163,10 +177,10 @@ def run_trace_cases( on_update: UpdateCallback, runner_args: Sequence[str] = (), ) -> tuple[int, HarnessRun]: - selected_scenarios: Final = frozenset(runner_args) + selected_scenarios, engine = runner_selection(runner_args) run: Final = HarnessRun.from_cases(cases) runnable_cases: Final = tuple(case for case in cases if isinstance(case.spec, ModuleCaseSpec)) - bridge_error: Final = ensure_trace_bridge(repo_root) if runnable_cases else None + bridge_error: Final = ensure_trace_bridge(repo_root) if runnable_cases and engine != "python" else None if bridge_error is not None: for harness_case in runnable_cases: _record_setup_failure(run, harness_case, bridge_error, "bridge") @@ -174,7 +188,7 @@ def run_trace_cases( on_update(run) return 1, run for harness_case in cases: - _run_case(run, harness_case, selected_scenarios, on_update) + _run_case(run, harness_case, selected_scenarios, on_update, engine) run.finished_at = monotonic() on_update(run) failed: Final = any( diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/chat_completions/case.py b/tests/rust-python-harness/strategies/trace_parity/sdk/chat_completions/case.py index 6be5afd60d6..1221f237570 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/chat_completions/case.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/chat_completions/case.py @@ -1,10 +1,15 @@ from __future__ import annotations -import json from typing import Final -from .....shared.parity.recorded_http import HttpHeader, RecordedHttpResponse from .....shared.tracing.steps import Engine, mapping +from ...fixtures import ( + anthropic_response_body, + anthropic_stream_events, + aws_event_stream_response, + json_response, + sse_response, +) from ...models import RouteFixture, RouteSpec, TraceScenario, TraceSuite COMMON_MAPPINGS: Final = ( @@ -24,6 +29,22 @@ COMMON_MAPPINGS: Final = ( mapping(rust_span="execute_chat_completions_provider_call"), mapping(rust_span="http_request", python_frame=r"AsyncHTTPHandler\.post$|HTTPHandler\.post$"), mapping(rust_span="transform_response", python_frame=r"(? RouteFixture: - response: Final = json.dumps( - { - "id": "msg_trace", - "type": "message", - "role": "assistant", - "model": "claude-sonnet-5", - "content": [{"type": "text", "text": "hello"}], - "stop_reason": "end_turn", - "stop_sequence": None, - "usage": {"input_tokens": 2, "output_tokens": 3}, - } - ).encode() return RouteFixture( kwargs={ "model": "anthropic/claude-sonnet-5", "messages": [{"role": "user", "content": "hello"}], **({"optional_params": {"max_tokens": 16}} if engine == "rust" else {"max_tokens": 16}), }, - provider_responses=( - RecordedHttpResponse.from_bytes( - 200, (HttpHeader(name="content-type", value="application/json"),), response - ), - ), + provider_responses=(json_response(anthropic_response_body()),), ) def _bedrock_fixture(engine: Engine, _base_url: str) -> RouteFixture: - response: Final = json.dumps( - { - "output": {"message": {"role": "assistant", "content": [{"text": "hello"}]}}, - "stopReason": "end_turn", - "usage": {"inputTokens": 2, "outputTokens": 3, "totalTokens": 5}, - "metrics": {"latencyMs": 1}, - } - ).encode() + response: Final[dict[str, object]] = { + "output": {"message": {"role": "assistant", "content": [{"text": "hello"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 2, "outputTokens": 3, "totalTokens": 5}, + "metrics": {"latencyMs": 1}, + } credentials: Final = { "aws_access_key_id": "test-access", "aws_secret_access_key": "test-secret", @@ -94,11 +97,60 @@ def _bedrock_fixture(engine: Engine, _base_url: str) -> RouteFixture: else {**credentials, "max_tokens": 16} ), }, + provider_responses=(json_response(response),), + ) + + +def _anthropic_stream_fixture(engine: Engine, _base_url: str) -> RouteFixture: + fixture: Final = _anthropic_fixture(engine, _base_url) + return fixture.derive( + kwargs={"stream": True}, + provider_responses=(sse_response(anthropic_stream_events()),), + consume_stream=True, + ) + + +def _bedrock_stream_fixture(engine: Engine, _base_url: str) -> RouteFixture: + fixture: Final = _bedrock_fixture(engine, _base_url) + events: Final[tuple[dict[str, object], ...]] = ( + {"messageStart": {"role": "assistant"}}, + {"contentBlockStart": {"contentBlockIndex": 0, "start": {}}}, + {"contentBlockDelta": {"contentBlockIndex": 0, "delta": {"text": "hello"}}}, + {"contentBlockStop": {"contentBlockIndex": 0}}, + {"messageStop": {"stopReason": "end_turn"}}, + {"metadata": {"usage": {"inputTokens": 2, "outputTokens": 1, "totalTokens": 3}}}, + ) + return fixture.derive( + kwargs={"stream": True}, + provider_responses=(aws_event_stream_response(events),), + consume_stream=True, + ) + + +def _provider_error_fixture(engine: Engine, _base_url: str) -> RouteFixture: + fixture: Final = _anthropic_fixture(engine, _base_url) + return fixture.derive( provider_responses=( - RecordedHttpResponse.from_bytes( - 200, (HttpHeader(name="content-type", value="application/json"),), response + json_response( + {"type": "error", "error": {"type": "invalid_request_error", "message": "bad request"}}, + status=400, ), ), + expected_failure=True, + ) + + +def _stream_error_fixture(engine: Engine, base_url: str) -> RouteFixture: + fixture: Final = _anthropic_fixture(engine, base_url) + events: Final = ( + anthropic_stream_events()[0], + ("error", {"type": "error", "error": {"type": "overloaded_error", "message": "overloaded"}}), + ) + return fixture.derive( + kwargs={"stream": True}, + provider_responses=(sse_response(events),), + expected_failure=True, + consume_stream=True, ) @@ -132,18 +184,64 @@ TRACE_SUITE: Final = TraceSuite( route=SPEC, scenarios=( TraceScenario( - name="anthropic", + name="sync-anthropic", fixture=_anthropic_fixture, - mappings=COMMON_MAPPINGS, - sync_mappings=SYNC_MAPPINGS, - async_mappings=ASYNC_MAPPINGS, + mappings=SYNC_MAPPINGS, + asynchronous=False, ), TraceScenario( - name="bedrock", + name="async-anthropic", + fixture=_anthropic_fixture, + mappings=ASYNC_MAPPINGS, + asynchronous=True, + ), + TraceScenario( + name="sync-anthropic-stream", + fixture=_anthropic_stream_fixture, + mappings=(*SYNC_MAPPINGS, *STREAM_MAPPINGS), + asynchronous=False, + ), + TraceScenario( + name="async-anthropic-stream", + fixture=_anthropic_stream_fixture, + mappings=(*ASYNC_MAPPINGS, *STREAM_MAPPINGS), + asynchronous=True, + ), + TraceScenario( + name="async-anthropic-provider-error", + fixture=_provider_error_fixture, + mappings=(*ASYNC_MAPPINGS, *FAILURE_MAPPINGS), + asynchronous=True, + ), + TraceScenario( + name="async-anthropic-stream-error", + fixture=_stream_error_fixture, + mappings=(*ASYNC_MAPPINGS, *STREAM_MAPPINGS, *FAILURE_MAPPINGS), + asynchronous=True, + ), + TraceScenario( + name="sync-bedrock", fixture=_bedrock_fixture, - mappings=BEDROCK_COMMON_MAPPINGS, - sync_mappings=BEDROCK_SYNC_MAPPINGS, - async_mappings=BEDROCK_ASYNC_MAPPINGS, + mappings=BEDROCK_SYNC_MAPPINGS, + asynchronous=False, + ), + TraceScenario( + name="async-bedrock", + fixture=_bedrock_fixture, + mappings=BEDROCK_ASYNC_MAPPINGS, + asynchronous=True, + ), + TraceScenario( + name="sync-bedrock-event-stream", + fixture=_bedrock_stream_fixture, + mappings=(*BEDROCK_SYNC_MAPPINGS, *STREAM_MAPPINGS), + asynchronous=False, + ), + TraceScenario( + name="async-bedrock-event-stream", + fixture=_bedrock_stream_fixture, + mappings=(*BEDROCK_ASYNC_MAPPINGS, *STREAM_MAPPINGS), + asynchronous=True, ), ), ) diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py b/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py index eb6c9233565..c249d09fa73 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py @@ -1,7 +1,7 @@ from __future__ import annotations import asyncio -from collections.abc import Awaitable +from collections.abc import AsyncIterable, Awaitable, Iterable from dataclasses import dataclass from pathlib import Path from typing import Final, Protocol, cast @@ -11,8 +11,8 @@ from ....shared.reporting.models import Surface from ....shared.tracing.native import TraceResponsePayload, native_trace_events from ....shared.tracing.profiler import FunctionTraceEvent, profile_python from ....shared.tracing.steps import Engine, pipeline_projection -from ..models import RouteFixture, RouteSpec, TraceExecutionFailure, TraceMode, TraceScenario -from ..reporting import TraceComparisonArtifact +from ..models import RouteFixture, RouteSpec, TraceEngine, TraceExecutionFailure, TraceScenario +from ..reporting import TraceArtifact class SdkCall(Protocol): @@ -25,10 +25,20 @@ class _CollectedTrace: error: str | None = None -def _invoke(function: SdkCall, kwargs: dict[str, object], *, asynchronous: bool) -> object: +def _invoke( + function: SdkCall, + kwargs: dict[str, object], + *, + asynchronous: bool, + consume_stream: bool = False, +) -> object: async def invoke_async() -> object: try: - return await cast(Awaitable[object], function(**kwargs)) + response: Final = await cast(Awaitable[object], function(**kwargs)) + if consume_stream and isinstance(response, AsyncIterable): + stream = cast(AsyncIterable[object], response) + return tuple([item async for item in stream]) + return response finally: await asyncio.sleep(0) from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER @@ -38,7 +48,10 @@ def _invoke(function: SdkCall, kwargs: dict[str, object], *, asynchronous: bool) if asynchronous: return asyncio.run(invoke_async()) - return function(**kwargs) + response: Final = function(**kwargs) + if consume_stream and isinstance(response, Iterable): + return tuple(cast(Iterable[object], response)) + return response def _entrypoint(spec: RouteSpec, engine: Engine, *, asynchronous: bool) -> SdkCall | TraceExecutionFailure: @@ -62,14 +75,6 @@ def _entrypoint(spec: RouteSpec, engine: Engine, *, asynchronous: bool) -> SdkCa return cast(SdkCall, getattr(owner, spec.python_entrypoints[int(asynchronous)])) -def _python_invocation_error(function: SdkCall, kwargs: dict[str, object], *, asynchronous: bool) -> str | None: - try: - _invoke(function, kwargs, asynchronous=asynchronous) - except Exception as error: - return f"{type(error).__name__}: {error}" - return None - - def _collect( function: SdkCall, fixture: RouteFixture, @@ -83,8 +88,19 @@ def _collect( return _CollectedTrace(native_trace_events(payload), payload.error) import litellm - with profile_python(Path(litellm.__file__).parent, threads=True) as profiler: - error: Final = _python_invocation_error(function, kwargs, asynchronous=asynchronous) + previous_suppress_debug_info: Final = litellm.suppress_debug_info + try: + if fixture.expected_failure: + litellm.suppress_debug_info = True + with profile_python(Path(litellm.__file__).parent, threads=True) as profiler: + error: str | None + try: + _invoke(function, kwargs, asynchronous=asynchronous, consume_stream=fixture.consume_stream) + error = None + except Exception as caught: + error = f"{type(caught).__name__}: {caught}" + finally: + litellm.suppress_debug_info = previous_suppress_debug_info return _CollectedTrace(tuple(profiler.events), error) @@ -108,6 +124,7 @@ def collect_trace( }, provider_responses=base_fixture.provider_responses, expected_failure=base_fixture.expected_failure, + consume_stream=base_fixture.consume_stream, ) collected: Final = _collect(function, fixture, engine, asynchronous=asynchronous) provider.take_requests(len(fixture.provider_responses)) @@ -129,18 +146,24 @@ def _failure_message(result: tuple[FunctionTraceEvent, ...] | TraceExecutionFail def execute_trace( - route: RouteSpec, scenario: TraceScenario, mode: TraceMode, surface: Surface -) -> TraceComparisonArtifact: - asynchronous: Final = mode == "async" - mappings: Final = scenario.mappings_for(mode) + route: RouteSpec, + scenario: TraceScenario, + surface: Surface, + engine: TraceEngine = "both", +) -> TraceArtifact: + mappings: Final = scenario.mappings scenario_route: Final = RouteSpec( route=route.route, python_entrypoints=route.python_entrypoints, rust_entrypoints=route.rust_entrypoints, fixture=scenario.fixture, ) - python_trace: Final = collect_trace(scenario_route, "python", asynchronous=asynchronous) - rust_trace: Final = collect_trace(scenario_route, "rust", asynchronous=asynchronous) + python_trace: Final = ( + collect_trace(scenario_route, "python", asynchronous=scenario.asynchronous) if engine != "rust" else () + ) + rust_trace: Final = ( + collect_trace(scenario_route, "rust", asynchronous=scenario.asynchronous) if engine != "python" else () + ) python_error: Final = _failure_message(python_trace) rust_error: Final = _failure_message(rust_trace) python_events: Final = python_trace if isinstance(python_trace, tuple) else () @@ -149,28 +172,22 @@ def execute_trace( python: Final = pipeline_projection("python", python_events, mappings) rust: Final = pipeline_projection("rust", rust_events, mappings) except ValueError as error: - return TraceComparisonArtifact.from_traces( + return TraceArtifact.from_traces( + engine=engine, surface=surface, sdk_function=route.route, scenario=scenario.name, - mode=mode, - mappings=mappings, - contract=scenario.contract, python=(), rust=(), - python_unmatched=0, python_error=f"harness: {error}", ) - return TraceComparisonArtifact.from_traces( + return TraceArtifact.from_traces( + engine=engine, surface=surface, sdk_function=route.route, scenario=scenario.name, - mode=mode, - mappings=mappings, - contract=scenario.contract, python=python.steps, rust=rust.steps, - python_unmatched=python.unmatched, python_error=python_error, rust_error=rust_error, ) diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/messages/case.py b/tests/rust-python-harness/strategies/trace_parity/sdk/messages/case.py index 27079c28cd8..211c454eadf 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/messages/case.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/messages/case.py @@ -1,14 +1,26 @@ from __future__ import annotations -import json from typing import Final -from .....shared.parity.recorded_http import HttpHeader, RecordedHttpResponse from .....shared.tracing.steps import Engine, mapping +from ...fixtures import ( + anthropic_response_body, + anthropic_stream_events, + aws_event_stream_response, + json_response, + sse_response, +) from ...models import RouteFixture, RouteSpec, TraceScenario, TraceSuite COMMON_MAPPINGS: Final = ( mapping(rust_span="messages", python_frame=r"anthropic_interface/messages/__init__\.py:\d+ a?create$"), + mapping(span="python_sanitize_empty_content", python_frame=r"strip_empty_content_blocks_from_anthropic_messages$"), + mapping(span="python_sanitize_tool_ids", python_frame=r"sanitize_tool_use_ids_in_anthropic_messages$"), + mapping( + span="python_flatten_web_search", python_frame=r"flatten_unencrypted_web_search_results_in_anthropic_messages$" + ), + mapping(span="python_cache_control", python_frame=r"AnthropicCacheControlHook\.maybe_inject_cache_control$"), + mapping(span="python_pre_request_hooks", python_frame=r"_execute_pre_request_hooks$"), mapping( span="python_messages_provider_config", python_frame=r"ProviderConfigManager\.get_provider_anthropic_messages_config$", @@ -30,10 +42,42 @@ COMMON_MAPPINGS: Final = ( ), mapping(rust_span="http_request", python_frame=r"AsyncHTTPHandler\.post$|HTTPHandler\.post$"), mapping(rust_span="transform_response", python_frame=r"(? RouteFixture: conversation: Final = {"messages": [{"role": "user", "content": "hello"}], "max_tokens": 16} - response: Final = json.dumps( - { - "id": "msg_trace", - "type": "message", - "role": "assistant", - "model": "claude-sonnet-5", - "content": [{"type": "text", "text": "hello"}], - "stop_reason": "end_turn", - "stop_sequence": None, - "usage": {"input_tokens": 2, "output_tokens": 3}, - } - ).encode() return RouteFixture( kwargs={ "model": f"{provider}/claude-sonnet-5", **({"body": {**conversation, "model": "claude-sonnet-5"}} if engine == "rust" else conversation), }, - provider_responses=( - RecordedHttpResponse.from_bytes( - 200, (HttpHeader(name="content-type", value="application/json"),), response - ), - ), + provider_responses=(json_response(anthropic_response_body()),), ) @@ -88,11 +156,170 @@ def _azure_fixture(engine: Engine, _base_url: str) -> RouteFixture: return _fixture(engine, "azure_ai") +def _bedrock_kwargs(engine: Engine) -> dict[str, object]: + conversation: Final = {"messages": [{"role": "user", "content": "hello"}], "max_tokens": 16} + return { + "model": "bedrock/anthropic.claude-3-sonnet-20240229-v1:0", + **( + {"body": {**conversation, "model": "anthropic.claude-3-sonnet-20240229-v1:0"}} + if engine == "rust" + else conversation + ), + "aws_access_key_id": "AKIAIOSFODNN7EXAMPLE", + "aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + "aws_region_name": "us-east-1", + } + + +def _bedrock_fixture(engine: Engine, _base_url: str) -> RouteFixture: + response_fixture: Final = _fixture(engine, "anthropic") + return RouteFixture(kwargs=_bedrock_kwargs(engine), provider_responses=response_fixture.provider_responses) + + +def _bedrock_retry_fixture(engine: Engine, _base_url: str) -> RouteFixture: + success_fixture: Final = _bedrock_fixture(engine, _base_url) + messages: Final = [ + {"role": "user", "content": "hello"}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "old reasoning", "signature": ""}, + {"type": "text", "text": "partial answer"}, + ], + }, + {"role": "user", "content": "continue"}, + ] + kwargs: Final = { + **_bedrock_kwargs(engine), + **( + {"body": {"messages": messages, "max_tokens": 16, "model": "anthropic.claude-3-sonnet-20240229-v1:0"}} + if engine == "rust" + else {"messages": messages} + ), + } + return success_fixture.derive( + kwargs=kwargs, + provider_responses=( + json_response({"message": "messages.1.content.0: Invalid `signature` in `thinking` block"}, status=400), + *success_fixture.provider_responses, + ), + ) + + +def _mock_fixture(engine: Engine, _base_url: str) -> RouteFixture: + fixture: Final = _fixture(engine, "anthropic") + return fixture.derive(kwargs={"mock_response": "hello from mock"}, provider_responses=()) + + +def _provider_error_fixture(engine: Engine, _base_url: str) -> RouteFixture: + fixture: Final = _fixture(engine, "anthropic") + return fixture.derive( + provider_responses=( + json_response( + {"type": "error", "error": {"type": "invalid_request_error", "message": "bad request"}}, + status=400, + ), + ), + expected_failure=True, + ) + + +def _sync_unsupported_fixture(engine: Engine, base_url: str) -> RouteFixture: + if engine == "rust": + return _anthropic_fixture(engine, base_url) + fixture: Final = _fixture(engine, "anthropic") + return fixture.derive(provider_responses=(), expected_failure=True) + + +def _stream_fixture_for(engine: Engine, provider: str) -> RouteFixture: + fixture: Final = _fixture(engine, provider) + return fixture.derive( + kwargs={"stream": True}, + provider_responses=(sse_response(anthropic_stream_events()),), + consume_stream=True, + ) + + +def _stream_fixture(engine: Engine, _base_url: str) -> RouteFixture: + return _stream_fixture_for(engine, "anthropic") + + +def _azure_stream_fixture(engine: Engine, _base_url: str) -> RouteFixture: + return _stream_fixture_for(engine, "azure_ai") + + +def _bedrock_stream_fixture(engine: Engine, base_url: str) -> RouteFixture: + fixture: Final = _bedrock_fixture(engine, base_url) + events: Final = tuple(payload for _, payload in anthropic_stream_events()) + return fixture.derive( + kwargs={"stream": True}, + provider_responses=(aws_event_stream_response(events),), + consume_stream=True, + ) + + +def _bedrock_stream_error_fixture(engine: Engine, base_url: str) -> RouteFixture: + fixture: Final = _bedrock_fixture(engine, base_url) + start: Final = anthropic_stream_events(model="anthropic.claude-3-sonnet-20240229-v1:0")[0][1] + return fixture.derive( + kwargs={"stream": True}, + provider_responses=(aws_event_stream_response((start, {"type": "message_stop"}), corrupt_last_frame=True),), + expected_failure=True, + consume_stream=True, + ) + + SPEC: Final = RouteSpec("messages", ("create", "acreate"), ("messages", "amessages"), _anthropic_fixture) TRACE_SUITE: Final = TraceSuite( route=SPEC, scenarios=( - TraceScenario(name="anthropic", fixture=_anthropic_fixture, mappings=ANTHROPIC_MAPPINGS, modes=("async",)), - TraceScenario(name="azure-ai", fixture=_azure_fixture, mappings=AZURE_MAPPINGS, modes=("async",)), + TraceScenario( + name="async-anthropic", fixture=_anthropic_fixture, mappings=ANTHROPIC_MAPPINGS, asynchronous=True + ), + TraceScenario(name="async-azure-ai", fixture=_azure_fixture, mappings=AZURE_MAPPINGS, asynchronous=True), + TraceScenario(name="async-bedrock", fixture=_bedrock_fixture, mappings=BEDROCK_MAPPINGS, asynchronous=True), + TraceScenario( + name="async-bedrock-invalid-thinking-retry", + fixture=_bedrock_retry_fixture, + mappings=RETRY_MAPPINGS, + asynchronous=True, + ), + TraceScenario(name="async-mock-response", fixture=_mock_fixture, mappings=MOCK_MAPPINGS, asynchronous=True), + TraceScenario( + name="async-anthropic-provider-error", + fixture=_provider_error_fixture, + mappings=ANTHROPIC_FAILURE_MAPPINGS, + asynchronous=True, + ), + TraceScenario( + name="async-anthropic-stream", + fixture=_stream_fixture, + mappings=(*ANTHROPIC_MAPPINGS, *STREAM_MAPPINGS), + asynchronous=True, + ), + TraceScenario( + name="async-azure-ai-stream", + fixture=_azure_stream_fixture, + mappings=(*AZURE_MAPPINGS, *STREAM_MAPPINGS), + asynchronous=True, + ), + TraceScenario( + name="async-bedrock-event-stream", + fixture=_bedrock_stream_fixture, + mappings=(*BEDROCK_MAPPINGS, *STREAM_MAPPINGS), + asynchronous=True, + ), + TraceScenario( + name="async-bedrock-event-stream-error", + fixture=_bedrock_stream_error_fixture, + mappings=(*BEDROCK_MAPPINGS, *STREAM_MAPPINGS, *FAILURE_MAPPINGS), + asynchronous=True, + ), + TraceScenario( + name="sync-unsupported", + fixture=_sync_unsupported_fixture, + mappings=ANTHROPIC_MAPPINGS, + asynchronous=False, + ), ), ) diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py b/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py index 2a4a1b3a152..effd1a0b4f6 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py @@ -358,53 +358,88 @@ TRACE_SUITE: Final = TraceSuite( route=SPEC, scenarios=( TraceScenario( - name="mistral", + name="sync-mistral", fixture=_mistral_fixture, - mappings=COMMON_MAPPINGS, - sync_mappings=(*SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), - async_mappings=(*ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + mappings=(*SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + asynchronous=False, ), TraceScenario( - name="mistral-callback-success", + name="async-mistral", + fixture=_mistral_fixture, + mappings=(*ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + asynchronous=True, + ), + TraceScenario( + name="sync-mistral-callback-success", fixture=_mistral_callback_success_fixture, - mappings=COMMON_MAPPINGS, - sync_mappings=CALLBACK_SUCCESS_SYNC_MAPPINGS, - async_mappings=CALLBACK_SUCCESS_ASYNC_MAPPINGS, + mappings=CALLBACK_SUCCESS_SYNC_MAPPINGS, + asynchronous=False, ), TraceScenario( - name="mistral-callback-failure", + name="async-mistral-callback-success", + fixture=_mistral_callback_success_fixture, + mappings=CALLBACK_SUCCESS_ASYNC_MAPPINGS, + asynchronous=True, + ), + TraceScenario( + name="sync-mistral-callback-failure", fixture=_mistral_callback_failure_fixture, - mappings=(*COMMON_MAPPINGS, FAILURE_CALLBACK_MAPPING), - sync_mappings=CALLBACK_FAILURE_SYNC_MAPPINGS, - async_mappings=CALLBACK_FAILURE_ASYNC_MAPPINGS, + mappings=CALLBACK_FAILURE_SYNC_MAPPINGS, + asynchronous=False, ), TraceScenario( - name="azure-ai", + name="async-mistral-callback-failure", + fixture=_mistral_callback_failure_fixture, + mappings=CALLBACK_FAILURE_ASYNC_MAPPINGS, + asynchronous=True, + ), + TraceScenario( + name="sync-azure-ai", fixture=_azure_fixture, - mappings=AZURE_COMMON_MAPPINGS, - sync_mappings=(*AZURE_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), - async_mappings=(*AZURE_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + mappings=(*AZURE_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + asynchronous=False, ), TraceScenario( - name="azure-document-intelligence", + name="async-azure-ai", + fixture=_azure_fixture, + mappings=(*AZURE_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + asynchronous=True, + ), + TraceScenario( + name="sync-azure-document-intelligence", fixture=_azure_document_intelligence_fixture, - mappings=DOCUMENT_INTELLIGENCE_COMMON_MAPPINGS, - sync_mappings=(*DOCUMENT_INTELLIGENCE_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), - async_mappings=(*DOCUMENT_INTELLIGENCE_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + mappings=(*DOCUMENT_INTELLIGENCE_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + asynchronous=False, ), TraceScenario( - name="vertex-ai", + name="async-azure-document-intelligence", + fixture=_azure_document_intelligence_fixture, + mappings=(*DOCUMENT_INTELLIGENCE_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + asynchronous=True, + ), + TraceScenario( + name="sync-vertex-ai", fixture=_vertex_fixture, - mappings=VERTEX_COMMON_MAPPINGS, - sync_mappings=(*VERTEX_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), - async_mappings=(*VERTEX_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + mappings=(*VERTEX_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + asynchronous=False, ), TraceScenario( - name="vertex-deepseek", + name="async-vertex-ai", + fixture=_vertex_fixture, + mappings=(*VERTEX_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + asynchronous=True, + ), + TraceScenario( + name="sync-vertex-deepseek", fixture=_vertex_deepseek_fixture, - mappings=DEEPSEEK_COMMON_MAPPINGS, - sync_mappings=(*DEEPSEEK_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), - async_mappings=(*DEEPSEEK_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + mappings=(*DEEPSEEK_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + asynchronous=False, + ), + TraceScenario( + name="async-vertex-deepseek", + fixture=_vertex_deepseek_fixture, + mappings=(*DEEPSEEK_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + asynchronous=True, ), ), ) diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/responses/case.py b/tests/rust-python-harness/strategies/trace_parity/sdk/responses/case.py new file mode 100644 index 00000000000..68d2eaba9bf --- /dev/null +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/responses/case.py @@ -0,0 +1,216 @@ +from __future__ import annotations + +from typing import Final + +from .....shared.tracing.steps import Engine, mapping +from ...fixtures import ( + anthropic_response_body, + anthropic_stream_events, + json_response, + responses_body, + responses_stream_events, + sse_response, +) +from ...models import RouteFixture, RouteSpec, TraceScenario, TraceSuite + +COMMON_MAPPINGS: Final = ( + mapping(span="python_responses", python_frame=r"responses/main\.py:\d+ a?responses$"), + mapping( + span="python_responses_provider_config", + python_frame=r"ProviderConfigManager\.get_provider_responses_api_config$", + ), + mapping(rust_span="responses_provider_config"), + mapping(rust_span="validate_environment", python_frame=r"validate_environment$"), + mapping(rust_span="complete_url", python_frame=r"get_complete_url$"), + mapping( + rust_span="transform_request", + python_frame=r"(? RouteFixture: + model: Final = "gpt-5" + return RouteFixture( + kwargs={ + "model": f"{provider}/{model}", + "input": "hello", + **({"body": {"model": model, "input": "hello"}} if engine == "rust" else {}), + }, + provider_responses=(json_response(responses_body(model=model)),), + ) + + +def _openai_fixture(engine: Engine, _base_url: str) -> RouteFixture: + return _native_fixture(engine, "openai") + + +def _azure_fixture(engine: Engine, _base_url: str) -> RouteFixture: + fixture: Final = _native_fixture(engine, "azure") + return fixture.derive(kwargs={"api_version": "2025-04-01-preview"}) + + +def _openai_stream_fixture(engine: Engine, _base_url: str) -> RouteFixture: + fixture: Final = _openai_fixture(engine, _base_url) + return fixture.derive( + kwargs={"stream": True}, + provider_responses=(sse_response(responses_stream_events()),), + consume_stream=True, + ) + + +def _provider_error_fixture(engine: Engine, _base_url: str) -> RouteFixture: + fixture: Final = _openai_fixture(engine, _base_url) + return fixture.derive( + provider_responses=( + json_response({"error": {"message": "bad request", "type": "invalid_request_error"}}, status=400), + ), + expected_failure=True, + ) + + +def _stream_failed_fixture(engine: Engine, base_url: str) -> RouteFixture: + fixture: Final = _openai_fixture(engine, base_url) + failed_response: Final[dict[str, object]] = { + **responses_body(), + "status": "failed", + "output": [], + "error": {"message": "stream failed", "type": "server_error", "code": "server_error"}, + } + events: Final = ( + ( + "response.created", + {"type": "response.created", "response": {**failed_response, "status": "in_progress", "error": None}}, + ), + ("response.failed", {"type": "response.failed", "response": failed_response}), + ) + return fixture.derive( + kwargs={"stream": True}, + provider_responses=(sse_response(events),), + expected_failure=True, + consume_stream=True, + ) + + +def _anthropic_bridge_fixture(engine: Engine, _base_url: str) -> RouteFixture: + return RouteFixture( + kwargs={ + "model": "anthropic/claude-sonnet-5", + "input": "hello", + "max_output_tokens": 16, + **({"body": {"model": "claude-sonnet-5", "input": "hello"}} if engine == "rust" else {}), + }, + provider_responses=(json_response(anthropic_response_body()),), + ) + + +def _anthropic_bridge_stream_fixture(engine: Engine, _base_url: str) -> RouteFixture: + fixture: Final = _anthropic_bridge_fixture(engine, _base_url) + return fixture.derive( + kwargs={"stream": True}, + provider_responses=(sse_response(anthropic_stream_events()),), + consume_stream=True, + ) + + +SPEC: Final = RouteSpec("responses", ("responses", "aresponses"), ("responses", "aresponses"), _openai_fixture) +TRACE_SUITE: Final = TraceSuite( + route=SPEC, + scenarios=( + TraceScenario(name="sync-openai", fixture=_openai_fixture, mappings=COMMON_MAPPINGS, asynchronous=False), + TraceScenario(name="async-openai", fixture=_openai_fixture, mappings=COMMON_MAPPINGS, asynchronous=True), + TraceScenario( + name="sync-openai-stream", + fixture=_openai_stream_fixture, + mappings=(*COMMON_MAPPINGS, *STREAM_MAPPINGS), + asynchronous=False, + ), + TraceScenario( + name="async-openai-stream", + fixture=_openai_stream_fixture, + mappings=(*COMMON_MAPPINGS, *STREAM_MAPPINGS), + asynchronous=True, + ), + TraceScenario( + name="async-openai-provider-error", + fixture=_provider_error_fixture, + mappings=(*COMMON_MAPPINGS, *FAILURE_MAPPINGS), + asynchronous=True, + ), + TraceScenario( + name="async-openai-stream-failed", + fixture=_stream_failed_fixture, + mappings=(*COMMON_MAPPINGS, *STREAM_MAPPINGS, *FAILURE_MAPPINGS), + asynchronous=True, + ), + TraceScenario(name="async-azure", fixture=_azure_fixture, mappings=AZURE_MAPPINGS, asynchronous=True), + TraceScenario( + name="async-anthropic-chat-bridge", + fixture=_anthropic_bridge_fixture, + mappings=BRIDGE_MAPPINGS, + asynchronous=True, + ), + TraceScenario( + name="async-anthropic-chat-bridge-stream", + fixture=_anthropic_bridge_stream_fixture, + mappings=( + *BRIDGE_MAPPINGS, + mapping(span="python_chat_stream_wrapper", python_frame=r"CustomStreamWrapper\.__init__$"), + mapping(span="python_chat_stream_next", python_frame=r"CustomStreamWrapper\.__anext__$"), + mapping( + span="python_responses_bridge_stream_iterator", + python_frame=r"LiteLLMCompletionStreamingIterator\.__init__$|LiteLLMCompletionStreamingIterator\.__anext__$", + ), + ), + asynchronous=True, + ), + ), +) diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/test_core_scenario_matrix.py b/tests/rust-python-harness/strategies/trace_parity/sdk/test_core_scenario_matrix.py new file mode 100644 index 00000000000..d0921398795 --- /dev/null +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/test_core_scenario_matrix.py @@ -0,0 +1,63 @@ +from __future__ import annotations + +from importlib import import_module +from typing import Final, cast + +from ..models import TraceSuite + + +def _suite(module: str) -> TraceSuite: + loaded: Final = import_module(module) + candidate: Final = cast(object, getattr(loaded, "TRACE_SUITE")) + assert isinstance(candidate, TraceSuite) + return candidate + + +def test_core_sdk_scenario_matrix_keeps_distinct_migration_paths() -> None: + chat: Final = _suite("tests.rust-python-harness.strategies.trace_parity.sdk.chat_completions.case") + messages: Final = _suite("tests.rust-python-harness.strategies.trace_parity.sdk.messages.case") + responses: Final = _suite("tests.rust-python-harness.strategies.trace_parity.sdk.responses.case") + + assert {(scenario.name, scenario.asynchronous) for scenario in chat.scenarios} >= { + ("sync-anthropic", False), + ("async-anthropic", True), + ("sync-anthropic-stream", False), + ("async-anthropic-stream", True), + ("async-anthropic-provider-error", True), + ("async-anthropic-stream-error", True), + ("sync-bedrock", False), + ("async-bedrock", True), + ("sync-bedrock-event-stream", False), + ("async-bedrock-event-stream", True), + } + assert {(scenario.name, scenario.asynchronous) for scenario in messages.scenarios} >= { + ("async-anthropic-stream", True), + ("async-azure-ai-stream", True), + ("async-bedrock-event-stream", True), + ("async-bedrock-event-stream-error", True), + ("async-bedrock-invalid-thinking-retry", True), + ("sync-unsupported", False), + } + assert {(scenario.name, scenario.asynchronous) for scenario in responses.scenarios} >= { + ("sync-openai", False), + ("async-openai", True), + ("sync-openai-stream", False), + ("async-openai-stream", True), + ("async-openai-provider-error", True), + ("async-openai-stream-failed", True), + ("async-azure", True), + ("async-anthropic-chat-bridge", True), + ("async-anthropic-chat-bridge-stream", True), + } + + +def test_core_gateway_matrix_keeps_downstream_streams_separate() -> None: + modules: Final = ( + "tests.rust-python-harness.strategies.trace_parity.gateway.chat_completions.case", + "tests.rust-python-harness.strategies.trace_parity.gateway.messages.case", + "tests.rust-python-harness.strategies.trace_parity.gateway.responses.case", + ) + + for module in modules: + suite = _suite(module) + assert any("downstream-stream" in scenario.name for scenario in suite.scenarios) diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/transcription/case.py b/tests/rust-python-harness/strategies/trace_parity/sdk/transcription/case.py index 2fb1054d10a..3b4d2e1447d 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/transcription/case.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/transcription/case.py @@ -92,11 +92,16 @@ TRACE_SUITE: Final = TraceSuite( route=SPEC, scenarios=( TraceScenario( - name="bedrock", + name="sync-bedrock", fixture=_fixture, - mappings=MAPPINGS, - sync_mappings=SYNC_MAPPINGS, - async_mappings=ASYNC_MAPPINGS, + mappings=SYNC_MAPPINGS, + asynchronous=False, + ), + TraceScenario( + name="async-bedrock", + fixture=_fixture, + mappings=ASYNC_MAPPINGS, + asynchronous=True, ), ), ) diff --git a/tests/rust-python-harness/strategies/trace_parity/test_reporting.py b/tests/rust-python-harness/strategies/trace_parity/test_reporting.py index 68264064c92..22cc87592b8 100644 --- a/tests/rust-python-harness/strategies/trace_parity/test_reporting.py +++ b/tests/rust-python-harness/strategies/trace_parity/test_reporting.py @@ -1,58 +1,46 @@ from __future__ import annotations -from collections.abc import Sequence from typing import Final, Literal import pytest from ...shared.reporting.models import CaseResult, Coverage, HarnessCase, ResultArtifact, RunStatus from ...shared.reporting.strategy import ModuleCaseSpec, NotImplementedCaseSpec -from ...shared.tracing.steps import PipelineStep, TraceContract, TraceMapping, mapping +from ...shared.tracing.steps import PipelineStep from . import reporting -from .reporting import TRACE_COMPARISON_ARTIFACT, TraceComparisonArtifact, render_trace_results - -MAPPINGS: Final = ( - mapping(rust_span="ocr", python_frame=r"ocr/main\.py:\d+ a?ocr$"), - mapping(rust_span="http_request", python_frame=r"AsyncHTTPHandler\.post$"), -) +from .reporting import TRACE_ARTIFACT, TraceArtifact, render_trace_results -def _result(comparison: TraceComparisonArtifact) -> CaseResult: +def _result(trace: TraceArtifact) -> CaseResult: case: Final = HarnessCase( strategy_id="trace_parity", strategy_label="Trace parity", - sdk_function=comparison.sdk_function, + sdk_function=trace.sdk_function, spec=ModuleCaseSpec(coverage=Coverage.PARTIAL, module="example"), - surface=comparison.surface, + surface=trace.surface, ) result: Final = CaseResult(case=case) - nodeid: Final = f"trace:sdk:{comparison.sdk_function}:{comparison.scenario}:{comparison.mode}" + nodeid: Final = f"trace:{trace.surface}:{trace.sdk_function}:{trace.scenario}" result.collected.add(nodeid) - result.record( - nodeid, - RunStatus.PASSED, - artifacts=(ResultArtifact(TRACE_COMPARISON_ARTIFACT, comparison.model_dump_json()),), - ) + result.record(nodeid, RunStatus.PASSED, artifacts=(ResultArtifact(TRACE_ARTIFACT, trace.model_dump_json()),)) return result -def _comparison( +def _trace( python: tuple[PipelineStep, ...], rust: tuple[PipelineStep, ...], *, - mappings: Sequence[TraceMapping] = MAPPINGS, rust_error: str | None = None, -) -> TraceComparisonArtifact: - return TraceComparisonArtifact.from_traces( + engine: Literal["python", "rust", "both"] = "both", + scenario: str = "sync-default", +) -> TraceArtifact: + return TraceArtifact.from_traces( + engine=engine, surface="sdk", sdk_function="ocr", - scenario="default", - mode="sync", - mappings=mappings, - contract=TraceContract(), + scenario=scenario, python=python, rust=rust, - python_unmatched=796, rust_error=rust_error, ) @@ -67,107 +55,55 @@ def _events(*items: tuple[str, int, str | None]) -> tuple[PipelineStep, ...]: return tuple(steps) -def test_renderer_shows_matching_python_and_rust_paths() -> None: - rust: Final = _events(("ocr", 0, None), ("http_request", 1, None)) +def test_renderer_prints_python_and_rust_traces_independently() -> None: python: Final = _events( ("ocr", 0, "ocr/main.py:88 aocr"), - ("http_request", 1, "http_handler.py:673 AsyncHTTPHandler.post"), + ("python_prepare", 1, "prep.py:1 python_prepare"), ) - - section: Final = render_trace_results((_result(_comparison(python, rust)),))[0] - report: Final = "\n\n".join(section.blocks) - - assert section.title == "SDK trace comparisons" - assert "Case: ocr" in report - assert "PYTHON (2 steps)\n1 aocr (ocr/main.py:88)\n2 AsyncHTTPHandler.post (http_handler.py:673)" in report - assert "RUST (2 steps)\nocr -> 1 aocr\n http_request -> 2 AsyncHTTPHandler.post" in report - assert "Mapping (identifier -> span)" not in report - assert "Trace: MATCH" in report - assert "Same steps, order, and nesting" in report - assert "Unseen mappings:" not in report - - -def test_renderer_reports_mappings_that_matched_nothing() -> None: - events: Final = _events(("ocr", 0, None)) - - section: Final = render_trace_results((_result(_comparison(events, events)),))[0] - report: Final = "\n\n".join(section.blocks) - - assert "Unseen mappings: http_request" in report - assert "Contract: FAIL" in report - - -def test_renderer_numbers_repeated_span_occurrences() -> None: - mappings: Final = (MAPPINGS[0], MAPPINGS[1]) - rust: Final = _events(("ocr", 0, None), ("http_request", 1, None), ("http_request", 1, None)) - python: Final = _events( - ("ocr", 0, "ocr/main.py:88 aocr"), - ("http_request", 1, "http_handler.py:673 AsyncHTTPHandler.post"), - ("http_request", 1, "http_handler.py:673 AsyncHTTPHandler.post"), - ) - - report: Final = "\n\n".join(render_trace_results((_result(_comparison(python, rust, mappings=mappings)),))[0].blocks) - - assert "http_request#2" in report - - -def test_renderer_accepts_declared_engine_specific_steps() -> None: - mappings: Final = ( - *MAPPINGS[:1], - mapping(span="python_prepare", python_frame=r"python_prepare$"), - mapping(rust_span="rust_prepare"), - ) - python: Final = _events(("ocr", 0, None), ("python_prepare", 1, "prep.py:1 python_prepare")) rust: Final = _events(("ocr", 0, None), ("rust_prepare", 1, None)) - section: Final = render_trace_results((_result(_comparison(python, rust, mappings=mappings)),))[0] + section: Final = render_trace_results((_result(_trace(python, rust)),))[0] report: Final = "\n\n".join(section.blocks) - assert "2 python_prepare (prep.py:1) [python only]" in report - assert "rust_prepare -> [rust only]" in report - assert "Trace: MATCH" in report - assert "Contract: PASS" in report + assert section.title == "SDK traces" + assert "PYTHON (2 steps)\n1 aocr (ocr/main.py:88)\n2 python_prepare (prep.py:1)" in report + assert "RUST (2 steps)\n1 ocr\n2 rust_prepare" in report + assert "python only" not in report + assert "rust only" not in report + assert " -> " not in report + assert "Trace: MATCH" not in report + assert "Trace: DRIFT" not in report + assert "Contract:" not in report -def test_unavailable_check_reports_mode_from_nodeid() -> None: - case: Final = HarnessCase( - strategy_id="trace_parity", - strategy_label="Trace parity", - sdk_function="ocr", - spec=ModuleCaseSpec(coverage=Coverage.PARTIAL, module="example"), - surface="sdk", - ) - result: Final = CaseResult(case=case) - result.collected.add("trace:sdk:ocr:default:sync") - result.record("trace:sdk:ocr:default:sync", RunStatus.ERROR) +@pytest.mark.parametrize( + ("engine", "present", "absent"), + (("python", "PYTHON (1 steps)", "RUST"), ("rust", "RUST (1 steps)", "PYTHON")), +) +def test_renderer_prints_only_selected_engine(engine: Literal["python", "rust"], present: str, absent: str) -> None: + events: Final = _events(("ocr", 0, None)) - section: Final = render_trace_results((result,))[0] - report: Final = "\n\n".join(section.blocks) + report: Final = "\n\n".join(render_trace_results((_result(_trace(events, events, engine=engine)),))[0].blocks) - assert "Case: ocr" in report - assert "Scenario: default / Mode: sync" in report - assert "Trace: NOT AVAILABLE\nTest outcome: error" in report - assert "unknown mode" not in report + assert present in report + assert absent not in report def test_renderer_keeps_collected_trace_when_one_engine_errors() -> None: - python: Final = _events( - ("ocr", 0, "ocr/main.py:88 aocr"), - ("http_request", 1, "http_handler.py:673 AsyncHTTPHandler.post"), + python: Final = _events(("ocr", 0, "ocr/main.py:88 aocr")) + + report: Final = "\n\n".join( + render_trace_results( + (_result(_trace(python, (), rust_error="rust: native Rust bridge must include the trace-parity feature")),) + )[0].blocks ) - section: Final = render_trace_results( - (_result(_comparison(python, (), rust_error="rust: native Rust bridge must include the trace-parity feature")),) - )[0] - report: Final = "\n\n".join(section.blocks) - - assert "PYTHON (2 steps)\n1 aocr (ocr/main.py:88) [python only]" in report + assert "PYTHON (1 steps)\n1 aocr (ocr/main.py:88)" in report assert "Rust error: rust: native Rust bridge must include the trace-parity feature" in report assert "hint: rebuild the native bridge with the trace-parity feature" in report - assert "Contract: FAIL" in report -def test_renderer_groups_all_modes_under_one_case_header() -> None: +def test_unavailable_trace_reports_scenario_from_nodeid() -> None: case: Final = HarnessCase( strategy_id="trace_parity", strategy_label="Trace parity", @@ -176,76 +112,55 @@ def test_renderer_groups_all_modes_under_one_case_header() -> None: surface="sdk", ) result: Final = CaseResult(case=case) - events: Final = _events(("ocr", 0, None)) - modes: Final[tuple[Literal["sync", "async"], ...]] = ("sync", "async") - for mode in modes: - nodeid = f"trace:sdk:ocr:default:{mode}" - result.collected.add(nodeid) - comparison = TraceComparisonArtifact.from_traces( - surface="sdk", - sdk_function="ocr", - scenario="default", - mode=mode, - mappings=MAPPINGS, - contract=TraceContract(), - python=events, - rust=events, - python_unmatched=0, - ) - result.record( - nodeid, - RunStatus.PASSED, - artifacts=(ResultArtifact(TRACE_COMPARISON_ARTIFACT, comparison.model_dump_json()),), - ) + result.collected.add("trace:sdk:ocr:async-error") + result.record("trace:sdk:ocr:async-error", RunStatus.ERROR) - section: Final = render_trace_results((result,))[0] + report: Final = "\n\n".join(render_trace_results((result,))[0].blocks) + + assert "Scenario: async-error" in report + assert "Trace: NOT AVAILABLE\nTest outcome: error" in report + + +def test_renderer_groups_scenarios_under_one_case_header() -> None: + result: Final = _result(_trace(_events(("ocr", 0, None)), (), scenario="sync-default")) + async_trace: Final = _trace((), _events(("ocr", 0, None)), scenario="async-default") + nodeid: Final = "trace:sdk:ocr:async-default" + result.collected.add(nodeid) + result.record(nodeid, RunStatus.PASSED, artifacts=(ResultArtifact(TRACE_ARTIFACT, async_trace.model_dump_json()),)) + + report: Final = render_trace_results((result,))[0].blocks[0] - assert len(section.blocks) == 1 - report: Final = section.blocks[0] assert report.count("Case: ocr") == 1 - assert "Scenario: default / Mode: sync" in report - assert "Scenario: default / Mode: async" in report + assert "Scenario: sync-default" in report + assert "Scenario: async-default" in report def test_renderer_colors_every_trace_line_in_a_terminal(monkeypatch: pytest.MonkeyPatch) -> None: - rust: Final = _events(("ocr", 0, None), ("http_request", 1, None)) - python: Final = _events( - ("ocr", 0, "ocr/main.py:88 aocr"), - ("http_request", 1, "http_handler.py:673 AsyncHTTPHandler.post"), - ) + events: Final = _events(("ocr", 0, "ocr/main.py:88 aocr")) monkeypatch.setattr(reporting.sys.stdout, "isatty", lambda: True) monkeypatch.delenv("NO_COLOR", raising=False) - section: Final = render_trace_results((_result(_comparison(python, rust)),))[0] - report: Final = "\n\n".join(section.blocks) + report: Final = "\n\n".join(render_trace_results((_result(_trace(events, events)),))[0].blocks) - assert "\033[36mPYTHON\033[0m (2 steps)" in report + assert "\033[36mPYTHON\033[0m (1 steps)" in report assert "\033[36m1 aocr (ocr/main.py:88)\033[0m" in report - assert "\033[33mRUST\033[0m (2 steps)" in report - assert "\033[33mocr\033[0m -> \033[36m1 aocr\033[0m" in report - assert "\033[33mhttp_request\033[0m -> \033[36m2 AsyncHTTPHandler.post\033[0m" in report + assert "\033[33mRUST\033[0m (1 steps)" in report + assert "\033[33m1 ocr\033[0m" in report -def test_renderer_groups_cases_and_unavailable_entries_by_surface() -> None: - events: Final = _events(("ocr", 0, None)) - gateway_results: Final = tuple( - CaseResult( - case=HarnessCase( - strategy_id="trace_parity", - strategy_label="Trace parity", - sdk_function=sdk_function, - spec=NotImplementedCaseSpec(reason=f"No {sdk_function} case is registered."), - surface="gateway", - ), - status=RunStatus.NOT_IMPLEMENTED, - ) - for sdk_function in ("ocr", "messages") +def test_renderer_groups_unavailable_entries_by_surface() -> None: + gateway_result: Final = CaseResult( + case=HarnessCase( + strategy_id="trace_parity", + strategy_label="Trace parity", + sdk_function="messages", + spec=NotImplementedCaseSpec(reason="No messages case is registered."), + surface="gateway", + ), + status=RunStatus.NOT_IMPLEMENTED, ) - sections: Final = render_trace_results((_result(_comparison(events, events)), *gateway_results)) + sections: Final = render_trace_results((_result(_trace((), ())), gateway_result)) - assert tuple(section.title for section in sections) == ("SDK trace comparisons", "GATEWAY trace comparisons") - gateway_report: Final = "\n\n".join(sections[1].blocks) - assert gateway_report.count("Not implemented") == 1 - assert "- ocr: No ocr case is registered." in gateway_report - assert "- messages: No messages case is registered." in gateway_report + assert tuple(section.title for section in sections) == ("SDK traces", "GATEWAY traces") + assert "- messages: No messages case is registered." in "\n\n".join(sections[1].blocks) diff --git a/tests/rust-python-harness/strategies/trace_parity/test_runner.py b/tests/rust-python-harness/strategies/trace_parity/test_runner.py index 0fcc5860ff2..4992854e21d 100644 --- a/tests/rust-python-harness/strategies/trace_parity/test_runner.py +++ b/tests/rust-python-harness/strategies/trace_parity/test_runner.py @@ -1,12 +1,20 @@ from __future__ import annotations -from typing import Final +import importlib +from pathlib import Path +from typing import Final, cast + +import pytest + +import litellm from ...shared.reporting.models import Coverage, HarnessCase, HarnessRun, RunStatus, SdkFunction, Surface from ...shared.reporting.strategy import ModuleCaseSpec -from ...shared.tracing.steps import Engine +from ...shared.tracing.steps import Engine, PipelineStep from .models import GatewayRouteSpec, RouteFixture, RouteSpec, TraceScenario, TraceSuite -from .runner import run_trace_mode, scenario_nodeids, validate_trace_suite +from .reporting import TraceArtifact +from .runner import run_trace_cases, run_trace_scenario, runner_selection, scenario_nodeids, validate_trace_suite +from .sdk.execution import execute_trace def _fixture(_engine: Engine, _base_url: str) -> RouteFixture: @@ -27,46 +35,98 @@ def test_scenario_filtering_and_occurrence_node_ids() -> None: suite: Final = TraceSuite( route=RouteSpec("ocr", ("ocr", "aocr"), ("ocr", "aocr"), _fixture), scenarios=( - TraceScenario("one", _fixture, (), modes=("sync", "async")), - TraceScenario("two", _fixture, (), modes=("async",)), + TraceScenario("sync-one", _fixture, (), asynchronous=False), + TraceScenario("async-one", _fixture, (), asynchronous=True), + TraceScenario("async-two", _fixture, (), asynchronous=True), ), ) case: Final = _case() - nodes: Final = scenario_nodeids(suite, case, frozenset({"two"})) + nodes: Final = scenario_nodeids(suite, case, frozenset({"async-two"})) - assert tuple(nodeid for _, _, nodeid in nodes) == ("trace:sdk:ocr:two:async",) + assert tuple(nodeid for _, nodeid in nodes) == ("trace:sdk:ocr:async-two",) + + +def test_python_engine_is_separate_from_scenario_selection() -> None: + assert runner_selection(("mistral", "--engine=python")) == (frozenset({"mistral"}), "python") + + +def test_python_engine_skips_native_bridge(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + runner: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.runner") + case: Final = _case() + selected: list[tuple[frozenset[str], str]] = [] + + def reject_bridge(_repo_root: Path) -> str | None: + raise AssertionError("Python-only tracing must not inspect or build the native bridge") + + def capture_case( + _run: HarnessRun, + _case: HarnessCase, + scenarios: frozenset[str], + _on_update: object, + engine: str, + ) -> None: + selected.append((scenarios, engine)) + + monkeypatch.setattr(runner, "ensure_trace_bridge", reject_bridge) + monkeypatch.setattr(runner, "_run_case", capture_case) + + exit_code, _ = run_trace_cases((case,), tmp_path, lambda _: None, ("mistral", "--engine=python")) + + assert exit_code == 0 + assert selected == [(frozenset({"mistral"}), "python")] + + +def test_expected_provider_failure_omits_feedback_banner( + capsys: pytest.CaptureFixture[str], monkeypatch: pytest.MonkeyPatch +) -> None: + loaded: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.sdk.responses.case") + suite: Final = cast(TraceSuite, loaded.TRACE_SUITE) + scenario: Final = next(item for item in suite.scenarios if item.name == "async-openai-provider-error") + monkeypatch.setattr(litellm, "suppress_debug_info", False) + assert isinstance(suite.route, RouteSpec) + + result: Final = execute_trace(suite.route, scenario, "sdk", engine="python") + + assert result.python_error is None + assert "Give Feedback / Get Help" not in capsys.readouterr().out + assert litellm.suppress_debug_info is False def test_scenario_validation_rejects_duplicate_and_unsafe_names() -> None: route: Final = RouteSpec("ocr", ("ocr", "aocr"), ("ocr", "aocr"), _fixture) duplicate: Final = TraceSuite( route=route, - scenarios=(TraceScenario("same", _fixture, ()), TraceScenario("same", _fixture, ())), + scenarios=( + TraceScenario("sync-same", _fixture, (), asynchronous=False), + TraceScenario("sync-same", _fixture, (), asynchronous=False), + ), + ) + unsafe: Final = TraceSuite( + route=route, scenarios=(TraceScenario("sync-bad:name", _fixture, (), asynchronous=False),) ) - unsafe: Final = TraceSuite(route=route, scenarios=(TraceScenario("bad:name", _fixture, ()),)) case: Final = _case() assert validate_trace_suite(duplicate, case) is not None assert validate_trace_suite(unsafe, case) is not None -def test_scenario_validation_rejects_invalid_modes_and_route_registration() -> None: - invalid_modes: Final = TraceSuite( +def test_scenario_validation_rejects_invalid_names_and_route_registration() -> None: + invalid_name: Final = TraceSuite( route=RouteSpec("ocr", ("ocr", "aocr"), ("ocr", "aocr"), _fixture), - scenarios=(TraceScenario("invalid", _fixture, (), modes=("sync", "sync")),), + scenarios=(TraceScenario("bedrock", _fixture, (), asynchronous=True),), ) wrong_function: Final = TraceSuite( route=RouteSpec("messages", ("create", "acreate"), ("messages", "amessages"), _fixture), - scenarios=(TraceScenario("one", _fixture, ()),), + scenarios=(TraceScenario("sync-one", _fixture, (), asynchronous=False),), ) wrong_surface: Final = TraceSuite( route=GatewayRouteSpec("ocr"), - scenarios=(TraceScenario("one", _fixture, ()),), + scenarios=(TraceScenario("sync-one", _fixture, (), asynchronous=False),), ) case: Final = _case() - assert "unique sync/async modes" in (validate_trace_suite(invalid_modes, case) or "") + assert "start with sync- or async-" in (validate_trace_suite(invalid_name, case) or "") assert "does not match case function" in (validate_trace_suite(wrong_function, case) or "") assert "must use RouteSpec" in (validate_trace_suite(wrong_surface, case) or "") @@ -77,11 +137,35 @@ def test_invalid_route_dispatch_records_harness_error() -> None: result: Final = run.results[case.key] suite: Final = TraceSuite( route=GatewayRouteSpec("ocr"), - scenarios=(TraceScenario("one", _fixture, (), modes=("sync",)),), + scenarios=(TraceScenario("sync-one", _fixture, (), asynchronous=False),), ) - nodeid: Final = "trace:sdk:ocr:one:sync" + nodeid: Final = "trace:sdk:ocr:sync-one" - run_trace_mode(run, result, suite, suite.scenarios[0], "sync", "sdk", nodeid, lambda _: None) + run_trace_scenario(run, result, suite, suite.scenarios[0], "sdk", nodeid, lambda _: None) assert result.outcomes[nodeid] is RunStatus.ERROR assert run.failures == [(nodeid, "gateway route cannot run on the sdk surface")] + + +def test_different_python_and_rust_traces_pass(monkeypatch: pytest.MonkeyPatch) -> None: + runner: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.runner") + case: Final = _case() + run: Final = HarnessRun.from_cases((case,)) + result: Final = run.results[case.key] + suite: Final = TraceSuite( + route=RouteSpec("ocr", ("ocr", "aocr"), ("ocr", "aocr"), _fixture), + scenarios=(TraceScenario("sync-one", _fixture, (), asynchronous=False),), + ) + trace: Final = TraceArtifact.from_traces( + surface="sdk", + sdk_function="ocr", + scenario="sync-one", + python=(PipelineStep(0, None, "python_step", "python.py:1 python_step"),), + rust=(PipelineStep(0, None, "rust_step", "rust_step"),), + ) + monkeypatch.setattr(runner, "_execute_scenario", lambda *_args: trace) + + run_trace_scenario(run, result, suite, suite.scenarios[0], "sdk", "trace:sdk:ocr:sync-one", lambda _: None) + + assert result.outcomes["trace:sdk:ocr:sync-one"] is RunStatus.PASSED + assert run.failures == [] From eb2ffaae79654883228fd6a51b85039f69016882 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 14 Sep 2026 11:37:40 -0700 Subject: [PATCH 057/112] fix(harness): rebuild bridge for trace parity --- .../shared/native_build.py | 20 ++++++++++++++++++- .../shared/test_native_build.py | 11 +++++----- 2 files changed, 25 insertions(+), 6 deletions(-) diff --git a/tests/rust-python-harness/shared/native_build.py b/tests/rust-python-harness/shared/native_build.py index b50def98e3a..f67488cecb4 100644 --- a/tests/rust-python-harness/shared/native_build.py +++ b/tests/rust-python-harness/shared/native_build.py @@ -16,6 +16,11 @@ _RUST_ROOT: Final = "litellm-rust" _LOCKFILE: Final = "Cargo.lock" _SOURCE_SUFFIXES: Final = frozenset({".rs", ".toml"}) _FAILURE_OUTPUT_LINES: Final = 15 +_TRACE_CHECK: Final = ( + "from litellm.rust_bridge import get_native_bridge; " + "bridge = get_native_bridge(); " + "raise SystemExit(0 if bridge is not None and getattr(bridge, '_trace', None) is not None else 1)" +) def needs_rebuild(native_mtime: float | None, newest_source_mtime: float | None) -> bool: @@ -73,6 +78,17 @@ def _rebuild(repo_root: Path) -> tuple[bool, str]: return completed.returncode == 0, "\n".join(lines[-_FAILURE_OUTPUT_LINES:]) +def _installed_bridge_has_trace(repo_root: Path) -> bool: + completed: Final = subprocess.run( + (sys.executable, "-c", _TRACE_CHECK), + cwd=repo_root, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + check=False, + ) + return completed.returncode == 0 + + def trace_bridge_error() -> str | None: bridge: Final = get_native_bridge() if bridge is None: @@ -85,7 +101,9 @@ def trace_bridge_error() -> str | None: def ensure_trace_bridge(repo_root: Path) -> str | None: native_path: Final = _native_module_path() native_mtime: Final = native_path.stat().st_mtime if native_path is not None and native_path.exists() else None - rebuild_required: Final = needs_rebuild(native_mtime, _newest_source_mtime(repo_root)) or trace_bridge_error() is not None + rebuild_required: Final = needs_rebuild( + native_mtime, _newest_source_mtime(repo_root) + ) or not _installed_bridge_has_trace(repo_root) if rebuild_required: print(f"Rebuilding native Rust bridge ({BRIDGE_FEATURE} feature)...", flush=True) succeeded: Final diff --git a/tests/rust-python-harness/shared/test_native_build.py b/tests/rust-python-harness/shared/test_native_build.py index 982f787aefa..dc08bc1a2b6 100644 --- a/tests/rust-python-harness/shared/test_native_build.py +++ b/tests/rust-python-harness/shared/test_native_build.py @@ -105,13 +105,14 @@ def test_ensure_trace_bridge_rebuilds_when_trace_feature_is_missing( state.rebuilt = True return True, "" + def fake_get_native_bridge() -> SimpleNamespace: + assert state.rebuilt + return SimpleNamespace(_trace=object()) + monkeypatch.setattr(native_build, "_native_module_path", lambda: native) monkeypatch.setattr(native_build, "_rebuild", fake_rebuild) - monkeypatch.setattr( - native_build, - "get_native_bridge", - lambda: SimpleNamespace(_trace=object() if state.rebuilt else None), - ) + monkeypatch.setattr(native_build, "_installed_bridge_has_trace", lambda repo_root: False) + monkeypatch.setattr(native_build, "get_native_bridge", fake_get_native_bridge) message: Final = native_build.ensure_trace_bridge(tmp_path) From d4869ba7108bef55063bc667fd7df02105b05a0b Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 14 Sep 2026 11:38:13 -0700 Subject: [PATCH 058/112] refactor: print rust and python traces independently --- .../unit_tests_mapping/mapping_validator.py | 46 +++++++++---------- 1 file changed, 22 insertions(+), 24 deletions(-) diff --git a/tests/rust-python-harness/strategies/unit_tests_mapping/mapping_validator.py b/tests/rust-python-harness/strategies/unit_tests_mapping/mapping_validator.py index 98ea0b02e68..9dd79e860e6 100644 --- a/tests/rust-python-harness/strategies/unit_tests_mapping/mapping_validator.py +++ b/tests/rust-python-harness/strategies/unit_tests_mapping/mapping_validator.py @@ -38,30 +38,28 @@ def _trace_functions( python_functions: Final[dict[str, PythonFunctionIdentity]] = {} rust_functions: Final[dict[str, RustFunctionIdentity]] = {} for scenario in suite.scenarios: - for mode in scenario.modes: - route: Final = RouteSpec( - route=suite.route.route, - python_entrypoints=suite.route.python_entrypoints, - rust_entrypoints=suite.route.rust_entrypoints, - fixture=scenario.fixture, - ) - python_trace: Final = collect_trace(route, "python", asynchronous=mode == "async") - rust_trace: Final = collect_trace(route, "rust", asynchronous=mode == "async") - if isinstance(python_trace, TraceExecutionFailure): - raise ValueError(f"Python trace discovery failed for {scenario.name}/{mode}: {python_trace.message}") - if isinstance(rust_trace, TraceExecutionFailure): - raise ValueError(f"Rust trace discovery failed for {scenario.name}/{mode}: {rust_trace.message}") - mappings: Final = scenario.mappings_for(mode) - python_projection: Final = pipeline_projection("python", python_trace, mappings) - rust_projection: Final = pipeline_projection("rust", rust_trace, mappings) - for step in python_projection.steps: - if step.span in spec.trace_spans: - function: Final = PythonFunctionIdentity.from_trace(step.raw) - python_functions[function.raw] = function - for step in rust_projection.steps: - if step.span in spec.trace_spans: - function: Final = RustFunctionIdentity.from_trace(step.raw) - rust_functions[step.raw] = function + route: Final = RouteSpec( + route=suite.route.route, + python_entrypoints=suite.route.python_entrypoints, + rust_entrypoints=suite.route.rust_entrypoints, + fixture=scenario.fixture, + ) + python_trace: Final = collect_trace(route, "python", asynchronous=scenario.asynchronous) + rust_trace: Final = collect_trace(route, "rust", asynchronous=scenario.asynchronous) + if isinstance(python_trace, TraceExecutionFailure): + raise ValueError(f"Python trace discovery failed for {scenario.name}: {python_trace.message}") + if isinstance(rust_trace, TraceExecutionFailure): + raise ValueError(f"Rust trace discovery failed for {scenario.name}: {rust_trace.message}") + python_projection: Final = pipeline_projection("python", python_trace, scenario.mappings) + rust_projection: Final = pipeline_projection("rust", rust_trace, scenario.mappings) + for step in python_projection.steps: + if step.span in spec.trace_spans: + function: Final = PythonFunctionIdentity.from_trace(step.raw) + python_functions[function.raw] = function + for step in rust_projection.steps: + if step.span in spec.trace_spans: + function: Final = RustFunctionIdentity.from_trace(step.raw) + rust_functions[step.raw] = function if not python_functions or not rust_functions: raise ValueError(f"Python trace discovery found no functions for spans: {', '.join(spec.trace_spans)}") return ( From 82ef6ea6abbbc0502d4107d6d75fe9dac5355eb4 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 18:53:29 +0000 Subject: [PATCH 059/112] feat(proxy): add general_settings.allowed_file_extensions for /v1/files Opt-in allowlist for upload filename extensions, checked before the existing blocked_file_extensions blocklist and mapped through the same upload validation failure path. None keeps today's behaviour, [] rejects every upload, matching is case-insensitive on both sides, and a filename with no extension is rejected when the allowlist is set. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 6 +- .../openai_files_endpoints/files_endpoints.py | 6 + .../general_upload_validation.py | 68 +++++++++--- litellm/proxy/proxy_server.py | 4 + .../test_files_endpoint.py | 103 ++++++++++++++++++ .../test_general_upload_validation.py | 53 +++++++++ .../proxy/proxy_server/test_proxy_config.py | 24 ++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 7 +- 8 files changed, 254 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 49e0247aad9..655b3cf414b 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2635,9 +2635,13 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="max file size in MB for /v1/files uploads, for any purpose, if a file is larger than this size it will be rejected before being forwarded to the provider", ) + allowed_file_extensions: tuple[str, ...] | None = Field( + None, + description="the only file extensions (e.g. ['.jsonl', '.pdf', '.txt']) accepted on /v1/files uploads, for any purpose, matched case-insensitively against the uploaded filename. Files with any other extension, or none, are rejected. An empty list rejects every upload. Unset means no allowlist is applied", + ) blocked_file_extensions: tuple[str, ...] | None = Field( None, - description="file extensions (e.g. ['.exe', '.sh']) rejected on /v1/files uploads, for any purpose, matched case-insensitively against the uploaded filename", + description="file extensions (e.g. ['.exe', '.sh']) rejected on /v1/files uploads, for any purpose, matched case-insensitively against the uploaded filename. Deprecated in favour of allowed_file_extensions; still enforced, after the allowlist, when set", ) max_response_size_mb: int | None = Field( None, diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 07cdc33e306..ae6e222a863 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -78,6 +78,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( ) from litellm.proxy.openai_files_endpoints.general_upload_validation import ( MB, + check_allowed_extension, check_blocked_extension, check_unsafe_filename, check_upload_file_size, @@ -473,6 +474,11 @@ async def create_file( if general_size_failure is not None: raise_upload_validation_failure(general_size_failure) + allowed_extensions: Final = coerce_optional_str_list_setting(general_settings.get("allowed_file_extensions")) + allowed_extension_failure: Final = check_allowed_extension(file.filename, allowed_extensions) + if allowed_extension_failure is not None: + raise_upload_validation_failure(allowed_extension_failure) + blocked_extensions: Final = coerce_optional_str_list_setting(general_settings.get("blocked_file_extensions")) blocked_extension_failure: Final = check_blocked_extension(file.filename, blocked_extensions) if blocked_extension_failure is not None: diff --git a/litellm/proxy/openai_files_endpoints/general_upload_validation.py b/litellm/proxy/openai_files_endpoints/general_upload_validation.py index 8c59a520272..22f34b8bd87 100644 --- a/litellm/proxy/openai_files_endpoints/general_upload_validation.py +++ b/litellm/proxy/openai_files_endpoints/general_upload_validation.py @@ -2,8 +2,8 @@ Upload validation applied to every purpose at POST /v1/files. batch_file_validation.py checks the JSONL shape of purpose="batch" uploads; this -module applies the same fast-fail-before-forwarding shape (size cap, blocked -extensions, path-traversal filenames) regardless of purpose. +module applies the same fast-fail-before-forwarding shape (size cap, allowed and +blocked extensions, path-traversal filenames) regardless of purpose. """ from dataclasses import dataclass @@ -31,10 +31,13 @@ def coerce_optional_int_setting(raw: object) -> int | None: raise TypeError(f"expected an integer, got {raw!r}") -def coerce_optional_str_list_setting(raw: object) -> tuple[str, ...]: - """A general_settings value declared as an optional list of strings, e.g. blocked_file_extensions.""" +def coerce_optional_str_list_setting(raw: object) -> tuple[str, ...] | None: + """A general_settings value declared as an optional list of strings, e.g. allowed_file_extensions. + + None (unset) and [] (set to nothing) are different answers for an allowlist, so both survive. + """ if raw is None: - return () + return None if not isinstance(raw, list) or not all(isinstance(item, str) for item in raw): raise TypeError(f"expected a list of strings, got {raw!r}") return tuple(raw) @@ -46,6 +49,11 @@ class UploadedFileTooLarge: limit_mb: int +@dataclass(frozen=True, slots=True) +class UploadedFileExtensionNotAllowed: + extension: str + + @dataclass(frozen=True, slots=True) class UploadedFileBlockedExtension: extension: str @@ -56,7 +64,9 @@ class UploadedFileUnsafeFilename: filename: str -UploadValidationFailure = UploadedFileTooLarge | UploadedFileBlockedExtension | UploadedFileUnsafeFilename +UploadValidationFailure = ( + UploadedFileTooLarge | UploadedFileExtensionNotAllowed | UploadedFileBlockedExtension | UploadedFileUnsafeFilename +) def _file_size_bytes(file_source: bytes | BinaryIO) -> int: @@ -81,19 +91,36 @@ def check_upload_file_size( return None +def _normalized_extension(filename: str | None) -> str: + if not filename: + return "" + try: + return Path(safe_filename(filename)).suffix.lower() + except ValueError: + return "" + + +def check_allowed_extension( + filename: str | None, + allowed_extensions: tuple[str, ...] | None, +) -> UploadedFileExtensionNotAllowed | None: + """None means the allowlist is not configured; an empty tuple means nothing is allowed.""" + if allowed_extensions is None: + return None + extension: Final = _normalized_extension(filename) + normalized_allowed: Final = frozenset(item.lower() for item in allowed_extensions) + if extension and extension in normalized_allowed: + return None + return UploadedFileExtensionNotAllowed(extension=extension) + + def check_blocked_extension( filename: str | None, - blocked_extensions: tuple[str, ...], + blocked_extensions: tuple[str, ...] | None, ) -> UploadedFileBlockedExtension | None: - if not blocked_extensions or not filename: + if not blocked_extensions: return None - try: - extension: Final = Path(safe_filename(filename)).suffix.lower() - except ValueError: - return None - # The uploaded name's extension is normalized above; blocked_extensions comes - # straight from config.yaml or the DB and is normalized here too, so a - # differently-cased entry (".EXE") still catches a lowercase upload. + extension: Final = _normalized_extension(filename) normalized_blocked: Final = frozenset(item.lower() for item in blocked_extensions) if extension and extension in normalized_blocked: return UploadedFileBlockedExtension(extension=extension) @@ -128,6 +155,17 @@ def raise_upload_validation_failure(failure: UploadValidationFailure) -> NoRetur param="file", code=413, ) + case UploadedFileExtensionNotAllowed(extension=extension): + raise ProxyException( + message=( + (f"File extension '{extension}'" if extension else "A file without an extension") + + " is not in this proxy's allowed_file_extensions setting. " + "The file was not forwarded to the provider." + ), + type="invalid_request_error", + param="file", + code=400, + ) case UploadedFileBlockedExtension(extension=extension): raise ProxyException( message=( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d9f8e04ebda..43d2a149758 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7065,6 +7065,9 @@ class ProxyConfig: if "max_file_size_mb" not in self._yaml_general_settings_keys: general_settings["max_file_size_mb"] = _general_settings.get("max_file_size_mb") + if "allowed_file_extensions" not in self._yaml_general_settings_keys: + general_settings["allowed_file_extensions"] = _general_settings.get("allowed_file_extensions") + if "blocked_file_extensions" not in self._yaml_general_settings_keys: general_settings["blocked_file_extensions"] = _general_settings.get("blocked_file_extensions") @@ -17036,6 +17039,7 @@ _GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES: Final[Mapping[str, str]] = MappingPro "max_request_size_mb": "Integer", "max_batch_file_size_mb": "Integer", "max_file_size_mb": "Integer", + "allowed_file_extensions": "List", "blocked_file_extensions": "List", "max_response_size_mb": "Integer", "proxy_config_reload_interval_seconds": "Integer", diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index b696d9ebe5f..bd3aafd7865 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -4703,6 +4703,109 @@ def test_create_file_blocked_extension_unset_allows_everything(monkeypatch, llm_ assert len(forwarded_calls) == 1 +@pytest.mark.parametrize( + "filename", + ["payload.exe", "notes.txt", "README"], + ids=["other_extension", "text_extension", "no_extension"], +) +def test_create_file_extension_outside_allowlist_rejected_before_forwarding( + monkeypatch, llm_router: Router, filename: str +): + import litellm.proxy.proxy_server as ps + + forwarded_calls = _setup_batch_upload_endpoint(monkeypatch, llm_router) + monkeypatch.setitem(ps.general_settings, "allowed_file_extensions", [".jsonl"]) + + try: + response = client.post( + "/v1/files", + files={"file": (filename, b"MZ\x90\x00", "application/octet-stream")}, + data={"purpose": "user_data"}, + headers={"Authorization": "Bearer test-key"}, + ) + finally: + _teardown_batch_upload_endpoint() + + assert response.status_code == 400, response.text + error = response.json()["error"] + assert error["type"] == "invalid_request_error" + assert error["param"] == "file" + assert "allowed_file_extensions" in error["message"] + assert forwarded_calls == [] + + +def test_create_file_allowed_extension_forwards_case_insensitively(monkeypatch, llm_router: Router): + import litellm.proxy.proxy_server as ps + + forwarded_calls = _setup_batch_upload_endpoint(monkeypatch, llm_router) + monkeypatch.setitem(ps.general_settings, "allowed_file_extensions", [".JSONL"]) + + try: + response = client.post( + "/v1/files", + files={"file": ("input.jsonl", b'{"custom_id": "1"}\n', "application/jsonl")}, + data={"purpose": "user_data"}, + headers={"Authorization": "Bearer test-key"}, + ) + finally: + _teardown_batch_upload_endpoint() + + assert response.status_code == 200, response.text + assert len(forwarded_calls) == 1 + + +def test_create_file_empty_allowlist_rejects_every_upload(monkeypatch, llm_router: Router): + import litellm.proxy.proxy_server as ps + + forwarded_calls = _setup_batch_upload_endpoint(monkeypatch, llm_router) + monkeypatch.setitem(ps.general_settings, "allowed_file_extensions", []) + + try: + response = client.post( + "/v1/files", + files={"file": ("input.jsonl", b'{"custom_id": "1"}\n', "application/jsonl")}, + data={"purpose": "user_data"}, + headers={"Authorization": "Bearer test-key"}, + ) + finally: + _teardown_batch_upload_endpoint() + + assert response.status_code == 400, response.text + assert "allowed_file_extensions" in response.json()["error"]["message"] + assert forwarded_calls == [] + + +def test_create_file_allowlist_runs_before_blocklist(monkeypatch, llm_router: Router): + """An extension in both lists is refused by the allowlist message, and the blocklist still holds on its own.""" + import litellm.proxy.proxy_server as ps + + forwarded_calls = _setup_batch_upload_endpoint(monkeypatch, llm_router) + monkeypatch.setitem(ps.general_settings, "allowed_file_extensions", [".jsonl"]) + monkeypatch.setitem(ps.general_settings, "blocked_file_extensions", [".exe", ".jsonl"]) + + try: + denied_by_allowlist = client.post( + "/v1/files", + files={"file": ("payload.exe", b"MZ\x90\x00", "application/octet-stream")}, + data={"purpose": "user_data"}, + headers={"Authorization": "Bearer test-key"}, + ) + denied_by_blocklist = client.post( + "/v1/files", + files={"file": ("input.jsonl", b'{"custom_id": "1"}\n', "application/jsonl")}, + data={"purpose": "user_data"}, + headers={"Authorization": "Bearer test-key"}, + ) + finally: + _teardown_batch_upload_endpoint() + + assert denied_by_allowlist.status_code == 400, denied_by_allowlist.text + assert "allowed_file_extensions" in denied_by_allowlist.json()["error"]["message"] + assert denied_by_blocklist.status_code == 400, denied_by_blocklist.text + assert "blocked_file_extensions" in denied_by_blocklist.json()["error"]["message"] + assert forwarded_calls == [] + + def test_create_file_path_traversal_filename_rejected_before_forwarding(monkeypatch, llm_router: Router): """A filename carrying a directory-traversal component must never reach storage or the provider.""" forwarded_calls = _setup_batch_upload_endpoint(monkeypatch, llm_router) diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_general_upload_validation.py b/tests/test_litellm/proxy/openai_files_endpoint/test_general_upload_validation.py index 9f7dd914e4a..f5b558f3334 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_general_upload_validation.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_general_upload_validation.py @@ -1,4 +1,5 @@ import io +from pathlib import Path import pytest @@ -6,11 +7,14 @@ from litellm.proxy._types import ProxyException from litellm.proxy.openai_files_endpoints.general_upload_validation import ( MB, UploadedFileBlockedExtension, + UploadedFileExtensionNotAllowed, UploadedFileTooLarge, UploadedFileUnsafeFilename, + check_allowed_extension, check_blocked_extension, check_unsafe_filename, check_upload_file_size, + coerce_optional_str_list_setting, raise_upload_validation_failure, ) @@ -79,6 +83,45 @@ def test_no_filename_skips_extension_check(): assert check_blocked_extension(None, (".exe",)) is None +def test_allowed_extension_passes(): + assert check_allowed_extension("batch.jsonl", (".jsonl", ".pdf")) is None + + +@pytest.mark.parametrize("filename", ["payload.exe", "notes.txt", "archive.tar.gz"]) +def test_extension_outside_allowlist_rejected(filename): + assert check_allowed_extension(filename, (".jsonl", ".pdf")) == UploadedFileExtensionNotAllowed( + extension=Path(filename).suffix + ) + + +def test_allowed_extension_match_is_case_insensitive_for_upload(): + assert check_allowed_extension("batch.JSONL", (".jsonl",)) is None + + +def test_allowed_extension_match_is_case_insensitive_for_configured_value(): + assert check_allowed_extension("batch.jsonl", (".JSONL",)) is None + + +@pytest.mark.parametrize("filename", ["README", "", None, "../../"]) +def test_no_extension_rejected_when_allowlist_set(filename): + """The allowlist grants by extension, so a name that yields none has nothing to be granted for.""" + assert check_allowed_extension(filename, (".jsonl",)) == UploadedFileExtensionNotAllowed(extension="") + + +def test_empty_allowlist_rejects_everything(): + assert check_allowed_extension("batch.jsonl", ()) == UploadedFileExtensionNotAllowed(extension=".jsonl") + + +def test_unset_allowlist_skips_check(): + assert check_allowed_extension("payload.exe", None) is None + + +def test_coerce_str_list_setting_keeps_unset_and_empty_distinct(): + assert coerce_optional_str_list_setting(None) is None + assert coerce_optional_str_list_setting([]) == () + assert coerce_optional_str_list_setting([".jsonl"]) == (".jsonl",) + + def test_path_traversal_filename_rejected(): assert check_unsafe_filename("../../etc/passwd") == UploadedFileUnsafeFilename(filename="../../etc/passwd") @@ -112,6 +155,16 @@ def test_ordinary_filenames_allowed(filename): "413", ("15.0 MB", "max_file_size_mb", "10 MB", "not forwarded"), ), + ( + UploadedFileExtensionNotAllowed(extension=".exe"), + "400", + (".exe", "allowed_file_extensions", "not forwarded"), + ), + ( + UploadedFileExtensionNotAllowed(extension=""), + "400", + ("without an extension", "allowed_file_extensions", "not forwarded"), + ), ( UploadedFileBlockedExtension(extension=".exe"), "400", diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index e109b650da7..d6851116324 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -3469,6 +3469,30 @@ async def test_ProxyConfig__update_general_settings_cleared_db_max_batch_file_si assert ps.general_settings.get("max_batch_file_size_mb") is None +@pytest.mark.asyncio +async def test_ProxyConfig__update_general_settings_applies_db_allowed_file_extensions(monkeypatch): + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + pc = ProxyConfig() + await pc._update_general_settings({"allowed_file_extensions": [".jsonl"]}) + from litellm.proxy import proxy_server as ps + + assert ps.general_settings.get("allowed_file_extensions") == [".jsonl"] + + +@pytest.mark.asyncio +async def test_ProxyConfig__update_general_settings_yaml_allowed_file_extensions_wins_over_db(monkeypatch): + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"allowed_file_extensions": [".pdf"]}, + ) + pc = ProxyConfig() + pc._yaml_general_settings_keys = {"allowed_file_extensions"} + await pc._update_general_settings({"allowed_file_extensions": [".jsonl"]}) + from litellm.proxy import proxy_server as ps + + assert ps.general_settings.get("allowed_file_extensions") == [".pdf"] + + @pytest.mark.asyncio async def test_ProxyConfig__update_general_settings_none_input_noop(): pc = ProxyConfig() diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 839aa52fa84..8c4ed5d7ff5 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -25809,6 +25809,11 @@ export interface components { * @description If True, lets keys address Responses API ids that this proxy did not issue (raw provider ids, or ids issued before response-id encryption was configured). Such an id carries no owner, so no ownership check can run on it; ids this proxy did issue keep full ownership enforcement. Off by default, in which case an unrecognized response id is rejected with 403 */ allow_unmanaged_response_ids?: boolean | null; + /** + * Allowed File Extensions + * @description the only file extensions (e.g. ['.jsonl', '.pdf', '.txt']) accepted on /v1/files uploads, for any purpose, matched case-insensitively against the uploaded filename. Files with any other extension, or none, are rejected. An empty list rejects every upload. Unset means no allowlist is applied + */ + allowed_file_extensions?: string[] | null; /** * Allowed Routes * @description Proxy API Endpoints you want users to be able to access @@ -25831,7 +25836,7 @@ export interface components { background_health_checks?: boolean | null; /** * Blocked File Extensions - * @description file extensions (e.g. ['.exe', '.sh']) rejected on /v1/files uploads, for any purpose, matched case-insensitively against the uploaded filename + * @description file extensions (e.g. ['.exe', '.sh']) rejected on /v1/files uploads, for any purpose, matched case-insensitively against the uploaded filename. Deprecated in favour of allowed_file_extensions; still enforced, after the allowlist, when set */ blocked_file_extensions?: string[] | null; /** From cee7215b24344a71ed9619b5911dcfa209a0f59e Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 13 Sep 2026 04:04:33 +0000 Subject: [PATCH 060/112] feat(model_info): opt-in field-level backfill from fallback generalization rules Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../fallback_generalizations.py | 31 ++++- ...odel_prices_and_context_window_backup.json | 4 + litellm/utils.py | 6 + model_prices_and_context_window.json | 4 + .../test_fallback_generalizations.py | 124 +++++++++++++++++- 5 files changed, 166 insertions(+), 3 deletions(-) diff --git a/litellm/litellm_core_utils/fallback_generalizations.py b/litellm/litellm_core_utils/fallback_generalizations.py index 7739fc82c77..a5360309a70 100644 --- a/litellm/litellm_core_utils/fallback_generalizations.py +++ b/litellm/litellm_core_utils/fallback_generalizations.py @@ -34,6 +34,11 @@ rules never mix the two and never use ``extends``. A rule whose Rules are only consulted after exact and case-insensitive lookups miss, so an exact cost-map entry always takes precedence over any rule. +Rules flagged with ``backfill_exact_entries: true`` also fill only keys missing +from an exact cost-map entry, while values already present on the entry win on +conflict. Only flagged capability rules participate in this backfill; routing +rules never do. + Patterns are matched case-insensitively with ``re.search`` and are not implicitly anchored: a rule must include ``^`` and ``$`` to bind to the whole model name, otherwise it matches as a substring. Keeping anchoring in the regex makes the rule @@ -57,6 +62,7 @@ PATTERN_FIELD: Final = "pattern" MODEL_INFO_FIELD: Final = "model_info" PROVIDER_KEY: Final = "litellm_provider" LEGACY_EXTENDS_FIELD: Final = "extends" +BACKFILL_FIELD: Final = "backfill_exact_entries" def _resolve_legacy_extends(rules: list) -> list: @@ -98,6 +104,7 @@ class _RoutingRule: class _CapabilityRule: pattern: re.Pattern model_info: dict + backfill_exact_entries: bool _CompiledRule = _RoutingRule | _CapabilityRule @@ -125,8 +132,9 @@ def _compile_rule(rule: object) -> tuple[_CompiledRule, ...]: e, ) return () + backfill: Final = rule.get(BACKFILL_FIELD) is True if PROVIDER_KEY not in model_info: - return (_CapabilityRule(pattern=compiled, model_info=model_info),) + return (_CapabilityRule(pattern=compiled, model_info=model_info, backfill_exact_entries=backfill),) provider: Final = model_info[PROVIDER_KEY] if not isinstance(provider, str): verbose_logger.warning( @@ -140,7 +148,7 @@ def _compile_rule(rule: object) -> tuple[_CompiledRule, ...]: return (_RoutingRule(pattern=compiled, provider=provider),) return ( _RoutingRule(pattern=compiled, provider=provider), - _CapabilityRule(pattern=compiled, model_info=model_info), + _CapabilityRule(pattern=compiled, model_info=model_info, backfill_exact_entries=backfill), ) @@ -151,6 +159,7 @@ class _FallbackGeneralizations: self.rules: list = [] self.routing_rules: tuple = () self.capability_rules: tuple = () + self.backfill_rules: tuple = () def set_rules(self, rules: list | None) -> None: installed: Final = rules if isinstance(rules, list) else [] @@ -158,6 +167,7 @@ class _FallbackGeneralizations: self.rules = installed self.routing_rules = tuple(rule for rule in compiled if isinstance(rule, _RoutingRule)) self.capability_rules = tuple(rule for rule in compiled if isinstance(rule, _CapabilityRule)) + self.backfill_rules = tuple(rule for rule in self.capability_rules if rule.backfill_exact_entries) def match_routing(self, model: str) -> str | None: if not model: @@ -175,6 +185,14 @@ class _FallbackGeneralizations: return None return {key: value for model_info in matched for key, value in model_info.items()} + def match_backfill(self, model: str) -> dict | None: + if not model: + return None + matched = tuple(rule.model_info for rule in self.backfill_rules if rule.pattern.search(model) is not None) + if not matched: + return None + return {key: value for model_info in matched for key, value in model_info.items()} + _registry: Final = _FallbackGeneralizations() @@ -210,3 +228,12 @@ def match_capability_generalizations(model: str) -> dict | None: capability rule matches. O(number of rules); only call once exact lookups have missed. """ return _registry.match_capabilities(model) + + +def match_backfill_generalizations(model: str) -> dict | None: + """Return the union of flagged capability rules matching ``model``. + + Later rules override earlier ones on key conflicts. Returns ``None`` when no + flagged rule matches. O(number of rules); only call once exact lookups have matched. + """ + return _registry.match_backfill(model) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 89d582db151..6ba3c240a94 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -57717,6 +57717,7 @@ { "name": "claude-adaptive-thinking", "pattern": "claude-[a-z]+-(?:4[-._](?:[6-9]|[1-9]\\d)(?!\\d)|(?:[5-9]|[1-9]\\d)(?!\\d)(?:[-._]\\d{1,2}(?!\\d))?)", + "backfill_exact_entries": true, "description": "Claude at version 4.6 or higher, in any id shape that contains claude--: minors 4.6 through 4.99, any later major-minor, and bare 5+ majors so a new family shaped like claude-fable-5 matches. Requiring the claude- prefix keeps non-Claude names such as team-sonnet-5-1 out. The minor is capped at two digits so an 8-digit date suffix such as claude-opus-4-20250514 is never read as a >= 4.6 minor. Turns on adaptive thinking for new versions and new families with no code change.", "model_info": { "supports_adaptive_thinking": true @@ -57725,6 +57726,7 @@ { "name": "claude-legacy-thinking", "pattern": "claude-[a-z]+-4[-._]6(?!\\d)", + "backfill_exact_entries": true, "description": "Claude at version 4.6 exactly, in any id shape that contains claude--4-6 (dotted and underscored minors included, dated releases such as claude-sonnet-4-6-20260219 too). The 4.6 family is adaptive-thinking yet still accepts legacy thinking.type=enabled with budget_tokens, so the caller's hard budget cap is forwarded verbatim instead of being rewritten to an uncapped output_config.effort. The lookahead keeps two-digit minors such as 4-60 from matching. 4.7+ and 5+ majors reject the legacy shape and stay on the adaptive translation.", "model_info": { "supports_legacy_thinking": true @@ -57741,6 +57743,7 @@ { "name": "claude-mid-conversation-system", "pattern": "claude-[a-z]+-(?:4[-._](?:[89]|[1-9]\\d)(?!\\d)|(?:[5-9]|[1-9]\\d)(?!\\d)(?:[-._]\\d{1,2}(?!\\d))?)", + "backfill_exact_entries": true, "description": "Claude at version 4.8 or higher, in any id shape that contains claude--: minors 4.8 through 4.99, any later major-minor, and bare 5+ majors so a new family like claude-fable-5 matches. Anthropic introduced mid-conversation system messages with Opus 4.8 and every newer Claude keeps them; 4.7 and below reject the system role inside messages.", "model_info": { "supports_mid_conversation_system": true @@ -57757,6 +57760,7 @@ { "name": "openai-reasoning-family-baseline", "pattern": "^(?!.*search-api)(?:[a-z0-9_.-]+/)*(?:ft:)?(?:o[1-9]\\d*(?![a-z0-9])|gpt-[5-9](?:\\.\\d+)?(?![0-9.])|(?:gpt-\\d+(?:\\.\\d+)?(?:-[a-z0-9]+)*-)?(?:codex|deep-research|chat-latest)(?![a-z0-9]))", + "backfill_exact_entries": true, "description": "OpenAI reasoning families by id shape, under any provider namespace and with an optional ft: prefix: the o-series (o1, o3-pro, o4-mini), gpt-5 through gpt-9 majors including dotted minors and suffixed variants (gpt-5.5-cyber, gpt-6-astra), and the codex, deep-research and chat-latest lines when standalone or on a gpt base. gpt-5-search-api is excluded because it is a search-only surface. Every model here is a reasoning model, and the Responses API drops the caller's reasoning param for any mapped OpenAI model whose info lacks supports_reasoning, so an id the registry has not named yet keeps its reasoning settings instead of silently losing them. Rules lose to exact entries. Carries no mode and no pricing, so cost stays on the standard unpriced behavior.", "model_info": { "supports_reasoning": true diff --git a/litellm/utils.py b/litellm/utils.py index 04139a124b6..87f270fa757 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -83,6 +83,7 @@ from litellm.constants import ( ) from litellm.litellm_core_utils.core_helpers import normalize_drop_params from litellm.litellm_core_utils.fallback_generalizations import ( + match_backfill_generalizations, match_capability_generalizations, ) from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload @@ -5819,6 +5820,11 @@ def _get_model_info_helper( ): _model_info = None + if _model_info is not None and key is not None: + backfill: Final = match_backfill_generalizations(key) + if backfill is not None: + _model_info = {**{k: v for k, v in backfill.items() if k not in _model_info}, **_model_info} + if _model_info is None: generalization: Final = _get_model_info_from_generalization( model=model, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 89d582db151..6ba3c240a94 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -57717,6 +57717,7 @@ { "name": "claude-adaptive-thinking", "pattern": "claude-[a-z]+-(?:4[-._](?:[6-9]|[1-9]\\d)(?!\\d)|(?:[5-9]|[1-9]\\d)(?!\\d)(?:[-._]\\d{1,2}(?!\\d))?)", + "backfill_exact_entries": true, "description": "Claude at version 4.6 or higher, in any id shape that contains claude--: minors 4.6 through 4.99, any later major-minor, and bare 5+ majors so a new family shaped like claude-fable-5 matches. Requiring the claude- prefix keeps non-Claude names such as team-sonnet-5-1 out. The minor is capped at two digits so an 8-digit date suffix such as claude-opus-4-20250514 is never read as a >= 4.6 minor. Turns on adaptive thinking for new versions and new families with no code change.", "model_info": { "supports_adaptive_thinking": true @@ -57725,6 +57726,7 @@ { "name": "claude-legacy-thinking", "pattern": "claude-[a-z]+-4[-._]6(?!\\d)", + "backfill_exact_entries": true, "description": "Claude at version 4.6 exactly, in any id shape that contains claude--4-6 (dotted and underscored minors included, dated releases such as claude-sonnet-4-6-20260219 too). The 4.6 family is adaptive-thinking yet still accepts legacy thinking.type=enabled with budget_tokens, so the caller's hard budget cap is forwarded verbatim instead of being rewritten to an uncapped output_config.effort. The lookahead keeps two-digit minors such as 4-60 from matching. 4.7+ and 5+ majors reject the legacy shape and stay on the adaptive translation.", "model_info": { "supports_legacy_thinking": true @@ -57741,6 +57743,7 @@ { "name": "claude-mid-conversation-system", "pattern": "claude-[a-z]+-(?:4[-._](?:[89]|[1-9]\\d)(?!\\d)|(?:[5-9]|[1-9]\\d)(?!\\d)(?:[-._]\\d{1,2}(?!\\d))?)", + "backfill_exact_entries": true, "description": "Claude at version 4.8 or higher, in any id shape that contains claude--: minors 4.8 through 4.99, any later major-minor, and bare 5+ majors so a new family like claude-fable-5 matches. Anthropic introduced mid-conversation system messages with Opus 4.8 and every newer Claude keeps them; 4.7 and below reject the system role inside messages.", "model_info": { "supports_mid_conversation_system": true @@ -57757,6 +57760,7 @@ { "name": "openai-reasoning-family-baseline", "pattern": "^(?!.*search-api)(?:[a-z0-9_.-]+/)*(?:ft:)?(?:o[1-9]\\d*(?![a-z0-9])|gpt-[5-9](?:\\.\\d+)?(?![0-9.])|(?:gpt-\\d+(?:\\.\\d+)?(?:-[a-z0-9]+)*-)?(?:codex|deep-research|chat-latest)(?![a-z0-9]))", + "backfill_exact_entries": true, "description": "OpenAI reasoning families by id shape, under any provider namespace and with an optional ft: prefix: the o-series (o1, o3-pro, o4-mini), gpt-5 through gpt-9 majors including dotted minors and suffixed variants (gpt-5.5-cyber, gpt-6-astra), and the codex, deep-research and chat-latest lines when standalone or on a gpt base. gpt-5-search-api is excluded because it is a search-only surface. Every model here is a reasoning model, and the Responses API drops the caller's reasoning param for any mapped OpenAI model whose info lacks supports_reasoning, so an id the registry has not named yet keeps its reasoning settings instead of silently losing them. Rules lose to exact entries. Carries no mode and no pricing, so cost stays on the standard unpriced behavior.", "model_info": { "supports_reasoning": true diff --git a/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py b/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py index b6e656f282a..0221b9bbffc 100644 --- a/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py +++ b/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py @@ -11,11 +11,11 @@ import logging import pytest - import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.fallback_generalizations import ( get_fallback_generalization_rules, + match_backfill_generalizations, match_capability_generalizations, match_routing_generalization, set_fallback_generalizations, @@ -116,6 +116,42 @@ def test_capability_union_is_last_wins_in_file_order(restore_generalizations): } +def test_backfill_requires_per_rule_opt_in(restore_generalizations): + restore_generalizations( + [ + {"name": "base", "pattern": r"^acme-", "model_info": {"supports_reasoning": True}}, + { + "name": "opt-in", + "pattern": r"^acme-", + "backfill_exact_entries": True, + "model_info": {"supports_vision": True}, + }, + ] + ) + assert match_backfill_generalizations("acme-1") == {"supports_vision": True} + assert match_capability_generalizations("acme-1") == { + "supports_reasoning": True, + "supports_vision": True, + } + + restore_generalizations( + [{"name": "base", "pattern": r"^acme-", "model_info": {"supports_reasoning": True}}] + ) + assert match_backfill_generalizations("acme-1") is None + + restore_generalizations( + [ + { + "name": "route", + "pattern": r"^acme-", + "backfill_exact_entries": True, + "model_info": {"litellm_provider": "openai"}, + } + ] + ) + assert match_backfill_generalizations("acme-1") is None + + def test_routing_rules_are_excluded_from_capability_results(restore_generalizations): restore_generalizations( [ @@ -299,6 +335,58 @@ def test_exact_entry_takes_precedence_over_rule(restore_generalizations): assert info["input_cost_per_token"] != 999.0 +def test_exact_entries_backfill_only_missing_fields(restore_generalizations, monkeypatch): + monkeypatch.setattr( + litellm, + "model_cost", + { + **litellm.model_cost, + "acme-full": { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "litellm_provider": "openai", + "mode": "chat", + "max_tokens": 7, + "supports_reasoning": False, + }, + "acme-bare": { + "input_cost_per_token": 3e-6, + "output_cost_per_token": 4e-6, + "litellm_provider": "openai", + "mode": "chat", + }, + }, + ) + restore_generalizations( + [ + { + "name": "acme-backfill", + "pattern": r"^acme-", + "backfill_exact_entries": True, + "model_info": {"supports_reasoning": True, "max_tokens": 5}, + } + ] + ) + litellm.get_model_info.cache_clear() + + full = litellm.get_model_info("acme-full", custom_llm_provider="openai") + assert full["supports_reasoning"] is False + assert full["max_tokens"] == 7 + + bare = litellm.get_model_info("acme-bare", custom_llm_provider="openai") + assert bare["supports_reasoning"] is True + assert bare["max_tokens"] == 5 + assert bare["input_cost_per_token"] == 3e-6 + assert bare["key"] == "acme-bare" + + restore_generalizations( + [{"name": "acme-backfill", "pattern": r"^acme-", "model_info": {"supports_reasoning": True, "max_tokens": 5}}] + ) + litellm.get_model_info.cache_clear() + unflagged = litellm.get_model_info("acme-bare", custom_llm_provider="openai") + assert unflagged.get("supports_reasoning") is None + + # --------------------------------------------------------------------------- # # Shipped rules (bundled cost map) # --------------------------------------------------------------------------- # @@ -605,6 +693,10 @@ def test_shipped_wandb_rule_loses_to_mapped_non_reasoning_entries(shipped_cost_m assert litellm.supports_reasoning(model=model, custom_llm_provider="wandb") is False, model +def test_shipped_wandb_rule_does_not_backfill_mapped_entries(shipped_cost_map): + assert match_backfill_generalizations("wandb/meta-llama/Llama-3.1-8B-Instruct") is None + + def test_shipped_wandb_rule_is_anchored_to_the_wandb_namespace(shipped_cost_map): """``^wandb/`` is anchored, so it cannot leak onto another provider's ids.""" assert match_capability_generalizations("wandb/some-new-model") == {"supports_reasoning": True} @@ -722,3 +814,33 @@ def test_shipped_openai_reasoning_rule_skips_non_reasoning_gpt_ids(shipped_cost_ def test_shipped_openai_reasoning_rule_loses_to_mapped_entries(shipped_cost_map): assert "gpt-5-search-api" in litellm.model_cost assert litellm.supports_reasoning(model="gpt-5-search-api", custom_llm_provider="openai") is False + + +@pytest.mark.parametrize( + "model,provider", + [ + ("azure/us/o1-2024-12-17", "azure"), + ("github_copilot/gpt-5", "github_copilot"), + ], +) +def test_shipped_openai_reasoning_rule_backfills_mapped_entries(shipped_cost_map, model, provider): + assert model in litellm.model_cost + raw_entry = litellm.model_cost[model] + assert "supports_reasoning" not in raw_entry + model_without_provider = model.removeprefix(f"{provider}/") + assert litellm.supports_reasoning(model=model_without_provider, custom_llm_provider=provider) is True + info = litellm.get_model_info(model=model_without_provider, custom_llm_provider=provider) + assert info["input_cost_per_token"] == raw_entry.get("input_cost_per_token", 0) + + +def test_shipped_claude_thinking_rules_backfill_without_family_limits(shipped_cost_map): + model = "perplexity/anthropic/claude-sonnet-4-6" + assert model in litellm.model_cost + raw_entry = litellm.model_cost[model] + assert "supports_adaptive_thinking" not in raw_entry + assert "max_input_tokens" not in raw_entry + + info = litellm.get_model_info(model="anthropic/claude-sonnet-4-6", custom_llm_provider="perplexity") + assert info["supports_adaptive_thinking"] is True + assert info["supports_legacy_thinking"] is True + assert info.get("max_input_tokens") is None From 42c708670d9977219ecbcec92bf2fc56cf1b377c Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 13 Sep 2026 04:22:41 +0000 Subject: [PATCH 061/112] fix(model_info): guard backfill by mode, drop provider key, tighten claude major regex Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../fallback_generalizations.py | 5 ++- ...odel_prices_and_context_window_backup.json | 8 ++-- litellm/utils.py | 3 +- model_prices_and_context_window.json | 8 ++-- .../test_fallback_generalizations.py | 44 +++++++++++++++++++ 5 files changed, 58 insertions(+), 10 deletions(-) diff --git a/litellm/litellm_core_utils/fallback_generalizations.py b/litellm/litellm_core_utils/fallback_generalizations.py index a5360309a70..8a029f9a642 100644 --- a/litellm/litellm_core_utils/fallback_generalizations.py +++ b/litellm/litellm_core_utils/fallback_generalizations.py @@ -191,7 +191,10 @@ class _FallbackGeneralizations: matched = tuple(rule.model_info for rule in self.backfill_rules if rule.pattern.search(model) is not None) if not matched: return None - return {key: value for model_info in matched for key, value in model_info.items()} + backfill: Final = { + key: value for model_info in matched for key, value in model_info.items() if key != PROVIDER_KEY + } + return backfill or None _registry: Final = _FallbackGeneralizations() diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 6ba3c240a94..a06688faa5b 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -57716,9 +57716,9 @@ }, { "name": "claude-adaptive-thinking", - "pattern": "claude-[a-z]+-(?:4[-._](?:[6-9]|[1-9]\\d)(?!\\d)|(?:[5-9]|[1-9]\\d)(?!\\d)(?:[-._]\\d{1,2}(?!\\d))?)", + "pattern": "claude-[a-z]+-(?:4[-._](?:[6-9]|[1-9]\\d)(?!\\d)|[5-9](?!\\d)(?:[-._]\\d{1,2}(?!\\d))?)", "backfill_exact_entries": true, - "description": "Claude at version 4.6 or higher, in any id shape that contains claude--: minors 4.6 through 4.99, any later major-minor, and bare 5+ majors so a new family shaped like claude-fable-5 matches. Requiring the claude- prefix keeps non-Claude names such as team-sonnet-5-1 out. The minor is capped at two digits so an 8-digit date suffix such as claude-opus-4-20250514 is never read as a >= 4.6 minor. Turns on adaptive thinking for new versions and new families with no code change.", + "description": "Claude at version 4.6 or higher, in any id shape that contains claude--: minors 4.6 through 4.99, any later major-minor, and bare 5+ majors so a new family shaped like claude-fable-5 matches. Requiring the claude- prefix keeps non-Claude names such as team-sonnet-5-1 out. The minor is capped at two digits so an 8-digit date suffix such as claude-opus-4-20250514 is never read as a >= 4.6 minor. Two-digit majors are deliberately not matched so ids like claude-opus-41 (4.1) are not read as major 41. Turns on adaptive thinking for new versions and new families with no code change.", "model_info": { "supports_adaptive_thinking": true } @@ -57742,9 +57742,9 @@ }, { "name": "claude-mid-conversation-system", - "pattern": "claude-[a-z]+-(?:4[-._](?:[89]|[1-9]\\d)(?!\\d)|(?:[5-9]|[1-9]\\d)(?!\\d)(?:[-._]\\d{1,2}(?!\\d))?)", + "pattern": "claude-[a-z]+-(?:4[-._](?:[89]|[1-9]\\d)(?!\\d)|[5-9](?!\\d)(?:[-._]\\d{1,2}(?!\\d))?)", "backfill_exact_entries": true, - "description": "Claude at version 4.8 or higher, in any id shape that contains claude--: minors 4.8 through 4.99, any later major-minor, and bare 5+ majors so a new family like claude-fable-5 matches. Anthropic introduced mid-conversation system messages with Opus 4.8 and every newer Claude keeps them; 4.7 and below reject the system role inside messages.", + "description": "Claude at version 4.8 or higher, in any id shape that contains claude--: minors 4.8 through 4.99, any later major-minor, and bare 5+ majors so a new family like claude-fable-5 matches. Two-digit majors are deliberately not matched so ids like claude-opus-41 (4.1) are not read as major 41. Anthropic introduced mid-conversation system messages with Opus 4.8 and every newer Claude keeps them; 4.7 and below reject the system role inside messages.", "model_info": { "supports_mid_conversation_system": true } diff --git a/litellm/utils.py b/litellm/utils.py index 87f270fa757..74e7001effe 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -255,6 +255,7 @@ from litellm.types.utils import ( ) _CALL_TYPE_ENUM_MAP: Final[dict] = {ct.value: ct for ct in CallTypes} +_BACKFILL_MODES: Final = frozenset({"chat", "responses"}) # +-----------------------------------------------+ # | | @@ -5820,7 +5821,7 @@ def _get_model_info_helper( ): _model_info = None - if _model_info is not None and key is not None: + if _model_info is not None and key is not None and _model_info.get("mode", "chat") in _BACKFILL_MODES: backfill: Final = match_backfill_generalizations(key) if backfill is not None: _model_info = {**{k: v for k, v in backfill.items() if k not in _model_info}, **_model_info} diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 6ba3c240a94..a06688faa5b 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -57716,9 +57716,9 @@ }, { "name": "claude-adaptive-thinking", - "pattern": "claude-[a-z]+-(?:4[-._](?:[6-9]|[1-9]\\d)(?!\\d)|(?:[5-9]|[1-9]\\d)(?!\\d)(?:[-._]\\d{1,2}(?!\\d))?)", + "pattern": "claude-[a-z]+-(?:4[-._](?:[6-9]|[1-9]\\d)(?!\\d)|[5-9](?!\\d)(?:[-._]\\d{1,2}(?!\\d))?)", "backfill_exact_entries": true, - "description": "Claude at version 4.6 or higher, in any id shape that contains claude--: minors 4.6 through 4.99, any later major-minor, and bare 5+ majors so a new family shaped like claude-fable-5 matches. Requiring the claude- prefix keeps non-Claude names such as team-sonnet-5-1 out. The minor is capped at two digits so an 8-digit date suffix such as claude-opus-4-20250514 is never read as a >= 4.6 minor. Turns on adaptive thinking for new versions and new families with no code change.", + "description": "Claude at version 4.6 or higher, in any id shape that contains claude--: minors 4.6 through 4.99, any later major-minor, and bare 5+ majors so a new family shaped like claude-fable-5 matches. Requiring the claude- prefix keeps non-Claude names such as team-sonnet-5-1 out. The minor is capped at two digits so an 8-digit date suffix such as claude-opus-4-20250514 is never read as a >= 4.6 minor. Two-digit majors are deliberately not matched so ids like claude-opus-41 (4.1) are not read as major 41. Turns on adaptive thinking for new versions and new families with no code change.", "model_info": { "supports_adaptive_thinking": true } @@ -57742,9 +57742,9 @@ }, { "name": "claude-mid-conversation-system", - "pattern": "claude-[a-z]+-(?:4[-._](?:[89]|[1-9]\\d)(?!\\d)|(?:[5-9]|[1-9]\\d)(?!\\d)(?:[-._]\\d{1,2}(?!\\d))?)", + "pattern": "claude-[a-z]+-(?:4[-._](?:[89]|[1-9]\\d)(?!\\d)|[5-9](?!\\d)(?:[-._]\\d{1,2}(?!\\d))?)", "backfill_exact_entries": true, - "description": "Claude at version 4.8 or higher, in any id shape that contains claude--: minors 4.8 through 4.99, any later major-minor, and bare 5+ majors so a new family like claude-fable-5 matches. Anthropic introduced mid-conversation system messages with Opus 4.8 and every newer Claude keeps them; 4.7 and below reject the system role inside messages.", + "description": "Claude at version 4.8 or higher, in any id shape that contains claude--: minors 4.8 through 4.99, any later major-minor, and bare 5+ majors so a new family like claude-fable-5 matches. Two-digit majors are deliberately not matched so ids like claude-opus-41 (4.1) are not read as major 41. Anthropic introduced mid-conversation system messages with Opus 4.8 and every newer Claude keeps them; 4.7 and below reject the system role inside messages.", "model_info": { "supports_mid_conversation_system": true } diff --git a/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py b/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py index 0221b9bbffc..036abedd90c 100644 --- a/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py +++ b/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py @@ -139,6 +139,18 @@ def test_backfill_requires_per_rule_opt_in(restore_generalizations): ) assert match_backfill_generalizations("acme-1") is None + restore_generalizations( + [ + { + "name": "mixed", + "pattern": r"^acme-", + "backfill_exact_entries": True, + "model_info": {"litellm_provider": "openai", "supports_vision": True}, + } + ] + ) + assert match_backfill_generalizations("acme-1") == {"supports_vision": True} + restore_generalizations( [ { @@ -355,6 +367,12 @@ def test_exact_entries_backfill_only_missing_fields(restore_generalizations, mon "litellm_provider": "openai", "mode": "chat", }, + "acme-image": { + "input_cost_per_token": 5e-6, + "output_cost_per_token": 6e-6, + "litellm_provider": "openai", + "mode": "image_generation", + }, }, ) restore_generalizations( @@ -379,6 +397,9 @@ def test_exact_entries_backfill_only_missing_fields(restore_generalizations, mon assert bare["input_cost_per_token"] == 3e-6 assert bare["key"] == "acme-bare" + image = litellm.get_model_info("acme-image", custom_llm_provider="openai") + assert image.get("supports_reasoning") is None + restore_generalizations( [{"name": "acme-backfill", "pattern": r"^acme-", "model_info": {"supports_reasoning": True, "max_tokens": 5}}] ) @@ -503,6 +524,18 @@ def test_shipped_version_boundaries(shipped_cost_map, model, provider, adaptive, assert info.get("supports_mid_conversation_system") is mid_conversation, model +def test_shipped_claude_version_regex_excludes_undelimited_41(shipped_cost_map): + unmatched = match_capability_generalizations("github_copilot/claude-opus-41") + assert unmatched is None or "supports_adaptive_thinking" not in unmatched + assert unmatched is None or "supports_mid_conversation_system" not in unmatched + + for model in ("claude-opus-5", "claude-sonnet-4-8"): + matched = match_capability_generalizations(model) + assert matched is not None + assert matched["supports_adaptive_thinking"] is True + assert matched["supports_mid_conversation_system"] is True + + def test_shipped_rules_cover_new_families_like_fable_at_5_plus(shipped_cost_map): """Both version gates accept any claude-- id at major 5 or higher, bare major or major-minor, so a new family shaped like claude-fable-5 gets adaptive @@ -833,6 +866,17 @@ def test_shipped_openai_reasoning_rule_backfills_mapped_entries(shipped_cost_map assert info["input_cost_per_token"] == raw_entry.get("input_cost_per_token", 0) +def test_shipped_openai_reasoning_rule_skips_non_text_modes(shipped_cost_map): + model = "gemini/deep-research-pro-preview-12-2025" + assert model in litellm.model_cost + raw_entry = litellm.model_cost[model] + assert "supports_reasoning" not in raw_entry + assert raw_entry["mode"] == "image_generation" + + info = litellm.get_model_info("deep-research-pro-preview-12-2025", custom_llm_provider="gemini") + assert info.get("supports_reasoning") is None + + def test_shipped_claude_thinking_rules_backfill_without_family_limits(shipped_cost_map): model = "perplexity/anthropic/claude-sonnet-4-6" assert model in litellm.model_cost From fb2057fde7871c657c929d07135cfe7d343c7325 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 13 Sep 2026 04:27:36 +0000 Subject: [PATCH 062/112] refactor(model_info): rename backfill_exact_entries to fill_missing_fields Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../fallback_generalizations.py | 32 +++++++++---------- ...odel_prices_and_context_window_backup.json | 8 ++--- litellm/utils.py | 11 ++++--- model_prices_and_context_window.json | 8 ++--- .../test_fallback_generalizations.py | 26 +++++++-------- 5 files changed, 44 insertions(+), 41 deletions(-) diff --git a/litellm/litellm_core_utils/fallback_generalizations.py b/litellm/litellm_core_utils/fallback_generalizations.py index 8a029f9a642..816a6c10bf1 100644 --- a/litellm/litellm_core_utils/fallback_generalizations.py +++ b/litellm/litellm_core_utils/fallback_generalizations.py @@ -34,10 +34,10 @@ rules never mix the two and never use ``extends``. A rule whose Rules are only consulted after exact and case-insensitive lookups miss, so an exact cost-map entry always takes precedence over any rule. -Rules flagged with ``backfill_exact_entries: true`` also fill only keys missing +Rules flagged with ``fill_missing_fields: true`` also fill only keys missing from an exact cost-map entry, while values already present on the entry win on -conflict. Only flagged capability rules participate in this backfill; routing -rules never do. +conflict. Only flagged capability rules participate in this fill; routing rules +never do. Patterns are matched case-insensitively with ``re.search`` and are not implicitly anchored: a rule must include ``^`` and ``$`` to bind to the whole model name, @@ -62,7 +62,7 @@ PATTERN_FIELD: Final = "pattern" MODEL_INFO_FIELD: Final = "model_info" PROVIDER_KEY: Final = "litellm_provider" LEGACY_EXTENDS_FIELD: Final = "extends" -BACKFILL_FIELD: Final = "backfill_exact_entries" +FILL_MISSING_FIELDS_FIELD: Final = "fill_missing_fields" def _resolve_legacy_extends(rules: list) -> list: @@ -104,7 +104,7 @@ class _RoutingRule: class _CapabilityRule: pattern: re.Pattern model_info: dict - backfill_exact_entries: bool + fill_missing_fields: bool _CompiledRule = _RoutingRule | _CapabilityRule @@ -132,9 +132,9 @@ def _compile_rule(rule: object) -> tuple[_CompiledRule, ...]: e, ) return () - backfill: Final = rule.get(BACKFILL_FIELD) is True + fill_missing_fields: Final = rule.get(FILL_MISSING_FIELDS_FIELD) is True if PROVIDER_KEY not in model_info: - return (_CapabilityRule(pattern=compiled, model_info=model_info, backfill_exact_entries=backfill),) + return (_CapabilityRule(pattern=compiled, model_info=model_info, fill_missing_fields=fill_missing_fields),) provider: Final = model_info[PROVIDER_KEY] if not isinstance(provider, str): verbose_logger.warning( @@ -148,7 +148,7 @@ def _compile_rule(rule: object) -> tuple[_CompiledRule, ...]: return (_RoutingRule(pattern=compiled, provider=provider),) return ( _RoutingRule(pattern=compiled, provider=provider), - _CapabilityRule(pattern=compiled, model_info=model_info, backfill_exact_entries=backfill), + _CapabilityRule(pattern=compiled, model_info=model_info, fill_missing_fields=fill_missing_fields), ) @@ -159,7 +159,7 @@ class _FallbackGeneralizations: self.rules: list = [] self.routing_rules: tuple = () self.capability_rules: tuple = () - self.backfill_rules: tuple = () + self.fill_missing_rules: tuple = () def set_rules(self, rules: list | None) -> None: installed: Final = rules if isinstance(rules, list) else [] @@ -167,7 +167,7 @@ class _FallbackGeneralizations: self.rules = installed self.routing_rules = tuple(rule for rule in compiled if isinstance(rule, _RoutingRule)) self.capability_rules = tuple(rule for rule in compiled if isinstance(rule, _CapabilityRule)) - self.backfill_rules = tuple(rule for rule in self.capability_rules if rule.backfill_exact_entries) + self.fill_missing_rules = tuple(rule for rule in self.capability_rules if rule.fill_missing_fields) def match_routing(self, model: str) -> str | None: if not model: @@ -185,16 +185,16 @@ class _FallbackGeneralizations: return None return {key: value for model_info in matched for key, value in model_info.items()} - def match_backfill(self, model: str) -> dict | None: + def match_fill_missing(self, model: str) -> dict | None: if not model: return None - matched = tuple(rule.model_info for rule in self.backfill_rules if rule.pattern.search(model) is not None) + matched = tuple(rule.model_info for rule in self.fill_missing_rules if rule.pattern.search(model) is not None) if not matched: return None - backfill: Final = { + fill_missing: Final = { key: value for model_info in matched for key, value in model_info.items() if key != PROVIDER_KEY } - return backfill or None + return fill_missing or None _registry: Final = _FallbackGeneralizations() @@ -233,10 +233,10 @@ def match_capability_generalizations(model: str) -> dict | None: return _registry.match_capabilities(model) -def match_backfill_generalizations(model: str) -> dict | None: +def match_fill_missing_generalizations(model: str) -> dict | None: """Return the union of flagged capability rules matching ``model``. Later rules override earlier ones on key conflicts. Returns ``None`` when no flagged rule matches. O(number of rules); only call once exact lookups have matched. """ - return _registry.match_backfill(model) + return _registry.match_fill_missing(model) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index a06688faa5b..44f52248b97 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -57717,7 +57717,7 @@ { "name": "claude-adaptive-thinking", "pattern": "claude-[a-z]+-(?:4[-._](?:[6-9]|[1-9]\\d)(?!\\d)|[5-9](?!\\d)(?:[-._]\\d{1,2}(?!\\d))?)", - "backfill_exact_entries": true, + "fill_missing_fields": true, "description": "Claude at version 4.6 or higher, in any id shape that contains claude--: minors 4.6 through 4.99, any later major-minor, and bare 5+ majors so a new family shaped like claude-fable-5 matches. Requiring the claude- prefix keeps non-Claude names such as team-sonnet-5-1 out. The minor is capped at two digits so an 8-digit date suffix such as claude-opus-4-20250514 is never read as a >= 4.6 minor. Two-digit majors are deliberately not matched so ids like claude-opus-41 (4.1) are not read as major 41. Turns on adaptive thinking for new versions and new families with no code change.", "model_info": { "supports_adaptive_thinking": true @@ -57726,7 +57726,7 @@ { "name": "claude-legacy-thinking", "pattern": "claude-[a-z]+-4[-._]6(?!\\d)", - "backfill_exact_entries": true, + "fill_missing_fields": true, "description": "Claude at version 4.6 exactly, in any id shape that contains claude--4-6 (dotted and underscored minors included, dated releases such as claude-sonnet-4-6-20260219 too). The 4.6 family is adaptive-thinking yet still accepts legacy thinking.type=enabled with budget_tokens, so the caller's hard budget cap is forwarded verbatim instead of being rewritten to an uncapped output_config.effort. The lookahead keeps two-digit minors such as 4-60 from matching. 4.7+ and 5+ majors reject the legacy shape and stay on the adaptive translation.", "model_info": { "supports_legacy_thinking": true @@ -57743,7 +57743,7 @@ { "name": "claude-mid-conversation-system", "pattern": "claude-[a-z]+-(?:4[-._](?:[89]|[1-9]\\d)(?!\\d)|[5-9](?!\\d)(?:[-._]\\d{1,2}(?!\\d))?)", - "backfill_exact_entries": true, + "fill_missing_fields": true, "description": "Claude at version 4.8 or higher, in any id shape that contains claude--: minors 4.8 through 4.99, any later major-minor, and bare 5+ majors so a new family like claude-fable-5 matches. Two-digit majors are deliberately not matched so ids like claude-opus-41 (4.1) are not read as major 41. Anthropic introduced mid-conversation system messages with Opus 4.8 and every newer Claude keeps them; 4.7 and below reject the system role inside messages.", "model_info": { "supports_mid_conversation_system": true @@ -57760,7 +57760,7 @@ { "name": "openai-reasoning-family-baseline", "pattern": "^(?!.*search-api)(?:[a-z0-9_.-]+/)*(?:ft:)?(?:o[1-9]\\d*(?![a-z0-9])|gpt-[5-9](?:\\.\\d+)?(?![0-9.])|(?:gpt-\\d+(?:\\.\\d+)?(?:-[a-z0-9]+)*-)?(?:codex|deep-research|chat-latest)(?![a-z0-9]))", - "backfill_exact_entries": true, + "fill_missing_fields": true, "description": "OpenAI reasoning families by id shape, under any provider namespace and with an optional ft: prefix: the o-series (o1, o3-pro, o4-mini), gpt-5 through gpt-9 majors including dotted minors and suffixed variants (gpt-5.5-cyber, gpt-6-astra), and the codex, deep-research and chat-latest lines when standalone or on a gpt base. gpt-5-search-api is excluded because it is a search-only surface. Every model here is a reasoning model, and the Responses API drops the caller's reasoning param for any mapped OpenAI model whose info lacks supports_reasoning, so an id the registry has not named yet keeps its reasoning settings instead of silently losing them. Rules lose to exact entries. Carries no mode and no pricing, so cost stays on the standard unpriced behavior.", "model_info": { "supports_reasoning": true diff --git a/litellm/utils.py b/litellm/utils.py index 74e7001effe..70fea3db4ec 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -83,8 +83,8 @@ from litellm.constants import ( ) from litellm.litellm_core_utils.core_helpers import normalize_drop_params from litellm.litellm_core_utils.fallback_generalizations import ( - match_backfill_generalizations, match_capability_generalizations, + match_fill_missing_generalizations, ) from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload @@ -5822,9 +5822,12 @@ def _get_model_info_helper( _model_info = None if _model_info is not None and key is not None and _model_info.get("mode", "chat") in _BACKFILL_MODES: - backfill: Final = match_backfill_generalizations(key) - if backfill is not None: - _model_info = {**{k: v for k, v in backfill.items() if k not in _model_info}, **_model_info} + fill_missing: Final = match_fill_missing_generalizations(key) + if fill_missing is not None: + _model_info = { + **{k: v for k, v in fill_missing.items() if k not in _model_info}, + **_model_info, + } if _model_info is None: generalization: Final = _get_model_info_from_generalization( diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index a06688faa5b..44f52248b97 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -57717,7 +57717,7 @@ { "name": "claude-adaptive-thinking", "pattern": "claude-[a-z]+-(?:4[-._](?:[6-9]|[1-9]\\d)(?!\\d)|[5-9](?!\\d)(?:[-._]\\d{1,2}(?!\\d))?)", - "backfill_exact_entries": true, + "fill_missing_fields": true, "description": "Claude at version 4.6 or higher, in any id shape that contains claude--: minors 4.6 through 4.99, any later major-minor, and bare 5+ majors so a new family shaped like claude-fable-5 matches. Requiring the claude- prefix keeps non-Claude names such as team-sonnet-5-1 out. The minor is capped at two digits so an 8-digit date suffix such as claude-opus-4-20250514 is never read as a >= 4.6 minor. Two-digit majors are deliberately not matched so ids like claude-opus-41 (4.1) are not read as major 41. Turns on adaptive thinking for new versions and new families with no code change.", "model_info": { "supports_adaptive_thinking": true @@ -57726,7 +57726,7 @@ { "name": "claude-legacy-thinking", "pattern": "claude-[a-z]+-4[-._]6(?!\\d)", - "backfill_exact_entries": true, + "fill_missing_fields": true, "description": "Claude at version 4.6 exactly, in any id shape that contains claude--4-6 (dotted and underscored minors included, dated releases such as claude-sonnet-4-6-20260219 too). The 4.6 family is adaptive-thinking yet still accepts legacy thinking.type=enabled with budget_tokens, so the caller's hard budget cap is forwarded verbatim instead of being rewritten to an uncapped output_config.effort. The lookahead keeps two-digit minors such as 4-60 from matching. 4.7+ and 5+ majors reject the legacy shape and stay on the adaptive translation.", "model_info": { "supports_legacy_thinking": true @@ -57743,7 +57743,7 @@ { "name": "claude-mid-conversation-system", "pattern": "claude-[a-z]+-(?:4[-._](?:[89]|[1-9]\\d)(?!\\d)|[5-9](?!\\d)(?:[-._]\\d{1,2}(?!\\d))?)", - "backfill_exact_entries": true, + "fill_missing_fields": true, "description": "Claude at version 4.8 or higher, in any id shape that contains claude--: minors 4.8 through 4.99, any later major-minor, and bare 5+ majors so a new family like claude-fable-5 matches. Two-digit majors are deliberately not matched so ids like claude-opus-41 (4.1) are not read as major 41. Anthropic introduced mid-conversation system messages with Opus 4.8 and every newer Claude keeps them; 4.7 and below reject the system role inside messages.", "model_info": { "supports_mid_conversation_system": true @@ -57760,7 +57760,7 @@ { "name": "openai-reasoning-family-baseline", "pattern": "^(?!.*search-api)(?:[a-z0-9_.-]+/)*(?:ft:)?(?:o[1-9]\\d*(?![a-z0-9])|gpt-[5-9](?:\\.\\d+)?(?![0-9.])|(?:gpt-\\d+(?:\\.\\d+)?(?:-[a-z0-9]+)*-)?(?:codex|deep-research|chat-latest)(?![a-z0-9]))", - "backfill_exact_entries": true, + "fill_missing_fields": true, "description": "OpenAI reasoning families by id shape, under any provider namespace and with an optional ft: prefix: the o-series (o1, o3-pro, o4-mini), gpt-5 through gpt-9 majors including dotted minors and suffixed variants (gpt-5.5-cyber, gpt-6-astra), and the codex, deep-research and chat-latest lines when standalone or on a gpt base. gpt-5-search-api is excluded because it is a search-only surface. Every model here is a reasoning model, and the Responses API drops the caller's reasoning param for any mapped OpenAI model whose info lacks supports_reasoning, so an id the registry has not named yet keeps its reasoning settings instead of silently losing them. Rules lose to exact entries. Carries no mode and no pricing, so cost stays on the standard unpriced behavior.", "model_info": { "supports_reasoning": true diff --git a/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py b/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py index 036abedd90c..eaecbd0b87e 100644 --- a/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py +++ b/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py @@ -15,8 +15,8 @@ import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.fallback_generalizations import ( get_fallback_generalization_rules, - match_backfill_generalizations, match_capability_generalizations, + match_fill_missing_generalizations, match_routing_generalization, set_fallback_generalizations, ) @@ -116,19 +116,19 @@ def test_capability_union_is_last_wins_in_file_order(restore_generalizations): } -def test_backfill_requires_per_rule_opt_in(restore_generalizations): +def test_fill_missing_requires_per_rule_opt_in(restore_generalizations): restore_generalizations( [ {"name": "base", "pattern": r"^acme-", "model_info": {"supports_reasoning": True}}, { "name": "opt-in", "pattern": r"^acme-", - "backfill_exact_entries": True, + "fill_missing_fields": True, "model_info": {"supports_vision": True}, }, ] ) - assert match_backfill_generalizations("acme-1") == {"supports_vision": True} + assert match_fill_missing_generalizations("acme-1") == {"supports_vision": True} assert match_capability_generalizations("acme-1") == { "supports_reasoning": True, "supports_vision": True, @@ -137,31 +137,31 @@ def test_backfill_requires_per_rule_opt_in(restore_generalizations): restore_generalizations( [{"name": "base", "pattern": r"^acme-", "model_info": {"supports_reasoning": True}}] ) - assert match_backfill_generalizations("acme-1") is None + assert match_fill_missing_generalizations("acme-1") is None restore_generalizations( [ { "name": "mixed", "pattern": r"^acme-", - "backfill_exact_entries": True, + "fill_missing_fields": True, "model_info": {"litellm_provider": "openai", "supports_vision": True}, } ] ) - assert match_backfill_generalizations("acme-1") == {"supports_vision": True} + assert match_fill_missing_generalizations("acme-1") == {"supports_vision": True} restore_generalizations( [ { "name": "route", "pattern": r"^acme-", - "backfill_exact_entries": True, + "fill_missing_fields": True, "model_info": {"litellm_provider": "openai"}, } ] ) - assert match_backfill_generalizations("acme-1") is None + assert match_fill_missing_generalizations("acme-1") is None def test_routing_rules_are_excluded_from_capability_results(restore_generalizations): @@ -347,7 +347,7 @@ def test_exact_entry_takes_precedence_over_rule(restore_generalizations): assert info["input_cost_per_token"] != 999.0 -def test_exact_entries_backfill_only_missing_fields(restore_generalizations, monkeypatch): +def test_exact_entries_fill_only_missing_fields(restore_generalizations, monkeypatch): monkeypatch.setattr( litellm, "model_cost", @@ -380,7 +380,7 @@ def test_exact_entries_backfill_only_missing_fields(restore_generalizations, mon { "name": "acme-backfill", "pattern": r"^acme-", - "backfill_exact_entries": True, + "fill_missing_fields": True, "model_info": {"supports_reasoning": True, "max_tokens": 5}, } ] @@ -726,8 +726,8 @@ def test_shipped_wandb_rule_loses_to_mapped_non_reasoning_entries(shipped_cost_m assert litellm.supports_reasoning(model=model, custom_llm_provider="wandb") is False, model -def test_shipped_wandb_rule_does_not_backfill_mapped_entries(shipped_cost_map): - assert match_backfill_generalizations("wandb/meta-llama/Llama-3.1-8B-Instruct") is None +def test_shipped_wandb_rule_does_not_fill_missing_mapped_entries(shipped_cost_map): + assert match_fill_missing_generalizations("wandb/meta-llama/Llama-3.1-8B-Instruct") is None def test_shipped_wandb_rule_is_anchored_to_the_wandb_namespace(shipped_cost_map): From 812bbee0b31db5908274448261209de2f7585739 Mon Sep 17 00:00:00 2001 From: kerry Date: Mon, 14 Sep 2026 17:00:41 +0000 Subject: [PATCH 063/112] fix(model_info): scope fill_missing backfill to the rule's providers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../fallback_generalizations.py | 66 ++++++++++++++----- ...odel_prices_and_context_window_backup.json | 8 +-- litellm/utils.py | 4 +- model_prices_and_context_window.json | 8 +-- .../test_fallback_generalizations.py | 62 +++++++++++++---- 5 files changed, 107 insertions(+), 41 deletions(-) diff --git a/litellm/litellm_core_utils/fallback_generalizations.py b/litellm/litellm_core_utils/fallback_generalizations.py index 816a6c10bf1..da75faa4136 100644 --- a/litellm/litellm_core_utils/fallback_generalizations.py +++ b/litellm/litellm_core_utils/fallback_generalizations.py @@ -34,10 +34,10 @@ rules never mix the two and never use ``extends``. A rule whose Rules are only consulted after exact and case-insensitive lookups miss, so an exact cost-map entry always takes precedence over any rule. -Rules flagged with ``fill_missing_fields: true`` also fill only keys missing -from an exact cost-map entry, while values already present on the entry win on -conflict. Only flagged capability rules participate in this fill; routing rules -never do. +Rules flagged with ``fill_missing_for_providers: [..]`` also fill only keys +missing from an exact cost-map entry when the entry's ``litellm_provider`` is +listed, while values already present on the entry win on conflict. Only flagged +capability rules participate in this fill; routing rules never do. Patterns are matched case-insensitively with ``re.search`` and are not implicitly anchored: a rule must include ``^`` and ``$`` to bind to the whole model name, @@ -62,7 +62,7 @@ PATTERN_FIELD: Final = "pattern" MODEL_INFO_FIELD: Final = "model_info" PROVIDER_KEY: Final = "litellm_provider" LEGACY_EXTENDS_FIELD: Final = "extends" -FILL_MISSING_FIELDS_FIELD: Final = "fill_missing_fields" +FILL_MISSING_FOR_PROVIDERS_FIELD: Final = "fill_missing_for_providers" def _resolve_legacy_extends(rules: list) -> list: @@ -104,7 +104,7 @@ class _RoutingRule: class _CapabilityRule: pattern: re.Pattern model_info: dict - fill_missing_fields: bool + fill_missing_for_providers: frozenset[str] _CompiledRule = _RoutingRule | _CapabilityRule @@ -132,9 +132,29 @@ def _compile_rule(rule: object) -> tuple[_CompiledRule, ...]: e, ) return () - fill_missing_fields: Final = rule.get(FILL_MISSING_FIELDS_FIELD) is True + if FILL_MISSING_FOR_PROVIDERS_FIELD not in rule: + fill_missing_for_providers: Final[frozenset[str]] = frozenset() + else: + raw_fill_missing_for_providers: Final = rule.get(FILL_MISSING_FOR_PROVIDERS_FIELD) + if not isinstance(raw_fill_missing_for_providers, (list, tuple)) or not all( + isinstance(provider, str) for provider in raw_fill_missing_for_providers + ): + verbose_logger.warning( + "LiteLLM: skipping malformed fallback generalization rule %s " + "('%s' must be a list of provider strings).", + rule.get(NAME_FIELD, pattern), + FILL_MISSING_FOR_PROVIDERS_FIELD, + ) + return () + fill_missing_for_providers = frozenset(raw_fill_missing_for_providers) if PROVIDER_KEY not in model_info: - return (_CapabilityRule(pattern=compiled, model_info=model_info, fill_missing_fields=fill_missing_fields),) + return ( + _CapabilityRule( + pattern=compiled, + model_info=model_info, + fill_missing_for_providers=fill_missing_for_providers, + ), + ) provider: Final = model_info[PROVIDER_KEY] if not isinstance(provider, str): verbose_logger.warning( @@ -148,7 +168,11 @@ def _compile_rule(rule: object) -> tuple[_CompiledRule, ...]: return (_RoutingRule(pattern=compiled, provider=provider),) return ( _RoutingRule(pattern=compiled, provider=provider), - _CapabilityRule(pattern=compiled, model_info=model_info, fill_missing_fields=fill_missing_fields), + _CapabilityRule( + pattern=compiled, + model_info=model_info, + fill_missing_for_providers=fill_missing_for_providers, + ), ) @@ -167,7 +191,7 @@ class _FallbackGeneralizations: self.rules = installed self.routing_rules = tuple(rule for rule in compiled if isinstance(rule, _RoutingRule)) self.capability_rules = tuple(rule for rule in compiled if isinstance(rule, _CapabilityRule)) - self.fill_missing_rules = tuple(rule for rule in self.capability_rules if rule.fill_missing_fields) + self.fill_missing_rules = tuple(rule for rule in self.capability_rules if rule.fill_missing_for_providers) def match_routing(self, model: str) -> str | None: if not model: @@ -185,10 +209,14 @@ class _FallbackGeneralizations: return None return {key: value for model_info in matched for key, value in model_info.items()} - def match_fill_missing(self, model: str) -> dict | None: - if not model: + def match_fill_missing(self, model: str, provider: str) -> dict | None: + if not model or not provider: return None - matched = tuple(rule.model_info for rule in self.fill_missing_rules if rule.pattern.search(model) is not None) + matched = tuple( + rule.model_info + for rule in self.fill_missing_rules + if provider in rule.fill_missing_for_providers and rule.pattern.search(model) is not None + ) if not matched: return None fill_missing: Final = { @@ -233,10 +261,12 @@ def match_capability_generalizations(model: str) -> dict | None: return _registry.match_capabilities(model) -def match_fill_missing_generalizations(model: str) -> dict | None: - """Return the union of flagged capability rules matching ``model``. +def match_fill_missing_generalizations(model: str, provider: str) -> dict | None: + """Return flagged capability rules matching ``model`` for ``provider``. - Later rules override earlier ones on key conflicts. Returns ``None`` when no - flagged rule matches. O(number of rules); only call once exact lookups have matched. + Later rules override earlier ones on key conflicts. Only rules listing + ``provider`` in ``fill_missing_for_providers`` contribute. Returns ``None`` + when no flagged rule matches. O(number of rules); only call once exact + lookups have matched. """ - return _registry.match_fill_missing(model) + return _registry.match_fill_missing(model, provider) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 44f52248b97..c9e85c28420 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -57717,7 +57717,7 @@ { "name": "claude-adaptive-thinking", "pattern": "claude-[a-z]+-(?:4[-._](?:[6-9]|[1-9]\\d)(?!\\d)|[5-9](?!\\d)(?:[-._]\\d{1,2}(?!\\d))?)", - "fill_missing_fields": true, + "fill_missing_for_providers": ["anthropic"], "description": "Claude at version 4.6 or higher, in any id shape that contains claude--: minors 4.6 through 4.99, any later major-minor, and bare 5+ majors so a new family shaped like claude-fable-5 matches. Requiring the claude- prefix keeps non-Claude names such as team-sonnet-5-1 out. The minor is capped at two digits so an 8-digit date suffix such as claude-opus-4-20250514 is never read as a >= 4.6 minor. Two-digit majors are deliberately not matched so ids like claude-opus-41 (4.1) are not read as major 41. Turns on adaptive thinking for new versions and new families with no code change.", "model_info": { "supports_adaptive_thinking": true @@ -57726,7 +57726,7 @@ { "name": "claude-legacy-thinking", "pattern": "claude-[a-z]+-4[-._]6(?!\\d)", - "fill_missing_fields": true, + "fill_missing_for_providers": ["anthropic"], "description": "Claude at version 4.6 exactly, in any id shape that contains claude--4-6 (dotted and underscored minors included, dated releases such as claude-sonnet-4-6-20260219 too). The 4.6 family is adaptive-thinking yet still accepts legacy thinking.type=enabled with budget_tokens, so the caller's hard budget cap is forwarded verbatim instead of being rewritten to an uncapped output_config.effort. The lookahead keeps two-digit minors such as 4-60 from matching. 4.7+ and 5+ majors reject the legacy shape and stay on the adaptive translation.", "model_info": { "supports_legacy_thinking": true @@ -57743,7 +57743,7 @@ { "name": "claude-mid-conversation-system", "pattern": "claude-[a-z]+-(?:4[-._](?:[89]|[1-9]\\d)(?!\\d)|[5-9](?!\\d)(?:[-._]\\d{1,2}(?!\\d))?)", - "fill_missing_fields": true, + "fill_missing_for_providers": ["anthropic"], "description": "Claude at version 4.8 or higher, in any id shape that contains claude--: minors 4.8 through 4.99, any later major-minor, and bare 5+ majors so a new family like claude-fable-5 matches. Two-digit majors are deliberately not matched so ids like claude-opus-41 (4.1) are not read as major 41. Anthropic introduced mid-conversation system messages with Opus 4.8 and every newer Claude keeps them; 4.7 and below reject the system role inside messages.", "model_info": { "supports_mid_conversation_system": true @@ -57760,7 +57760,7 @@ { "name": "openai-reasoning-family-baseline", "pattern": "^(?!.*search-api)(?:[a-z0-9_.-]+/)*(?:ft:)?(?:o[1-9]\\d*(?![a-z0-9])|gpt-[5-9](?:\\.\\d+)?(?![0-9.])|(?:gpt-\\d+(?:\\.\\d+)?(?:-[a-z0-9]+)*-)?(?:codex|deep-research|chat-latest)(?![a-z0-9]))", - "fill_missing_fields": true, + "fill_missing_for_providers": ["openai"], "description": "OpenAI reasoning families by id shape, under any provider namespace and with an optional ft: prefix: the o-series (o1, o3-pro, o4-mini), gpt-5 through gpt-9 majors including dotted minors and suffixed variants (gpt-5.5-cyber, gpt-6-astra), and the codex, deep-research and chat-latest lines when standalone or on a gpt base. gpt-5-search-api is excluded because it is a search-only surface. Every model here is a reasoning model, and the Responses API drops the caller's reasoning param for any mapped OpenAI model whose info lacks supports_reasoning, so an id the registry has not named yet keeps its reasoning settings instead of silently losing them. Rules lose to exact entries. Carries no mode and no pricing, so cost stays on the standard unpriced behavior.", "model_info": { "supports_reasoning": true diff --git a/litellm/utils.py b/litellm/utils.py index 70fea3db4ec..6af66010bd7 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5822,7 +5822,9 @@ def _get_model_info_helper( _model_info = None if _model_info is not None and key is not None and _model_info.get("mode", "chat") in _BACKFILL_MODES: - fill_missing: Final = match_fill_missing_generalizations(key) + fill_missing: Final = match_fill_missing_generalizations( + key, _model_info.get("litellm_provider", "") + ) if fill_missing is not None: _model_info = { **{k: v for k, v in fill_missing.items() if k not in _model_info}, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 44f52248b97..c9e85c28420 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -57717,7 +57717,7 @@ { "name": "claude-adaptive-thinking", "pattern": "claude-[a-z]+-(?:4[-._](?:[6-9]|[1-9]\\d)(?!\\d)|[5-9](?!\\d)(?:[-._]\\d{1,2}(?!\\d))?)", - "fill_missing_fields": true, + "fill_missing_for_providers": ["anthropic"], "description": "Claude at version 4.6 or higher, in any id shape that contains claude--: minors 4.6 through 4.99, any later major-minor, and bare 5+ majors so a new family shaped like claude-fable-5 matches. Requiring the claude- prefix keeps non-Claude names such as team-sonnet-5-1 out. The minor is capped at two digits so an 8-digit date suffix such as claude-opus-4-20250514 is never read as a >= 4.6 minor. Two-digit majors are deliberately not matched so ids like claude-opus-41 (4.1) are not read as major 41. Turns on adaptive thinking for new versions and new families with no code change.", "model_info": { "supports_adaptive_thinking": true @@ -57726,7 +57726,7 @@ { "name": "claude-legacy-thinking", "pattern": "claude-[a-z]+-4[-._]6(?!\\d)", - "fill_missing_fields": true, + "fill_missing_for_providers": ["anthropic"], "description": "Claude at version 4.6 exactly, in any id shape that contains claude--4-6 (dotted and underscored minors included, dated releases such as claude-sonnet-4-6-20260219 too). The 4.6 family is adaptive-thinking yet still accepts legacy thinking.type=enabled with budget_tokens, so the caller's hard budget cap is forwarded verbatim instead of being rewritten to an uncapped output_config.effort. The lookahead keeps two-digit minors such as 4-60 from matching. 4.7+ and 5+ majors reject the legacy shape and stay on the adaptive translation.", "model_info": { "supports_legacy_thinking": true @@ -57743,7 +57743,7 @@ { "name": "claude-mid-conversation-system", "pattern": "claude-[a-z]+-(?:4[-._](?:[89]|[1-9]\\d)(?!\\d)|[5-9](?!\\d)(?:[-._]\\d{1,2}(?!\\d))?)", - "fill_missing_fields": true, + "fill_missing_for_providers": ["anthropic"], "description": "Claude at version 4.8 or higher, in any id shape that contains claude--: minors 4.8 through 4.99, any later major-minor, and bare 5+ majors so a new family like claude-fable-5 matches. Two-digit majors are deliberately not matched so ids like claude-opus-41 (4.1) are not read as major 41. Anthropic introduced mid-conversation system messages with Opus 4.8 and every newer Claude keeps them; 4.7 and below reject the system role inside messages.", "model_info": { "supports_mid_conversation_system": true @@ -57760,7 +57760,7 @@ { "name": "openai-reasoning-family-baseline", "pattern": "^(?!.*search-api)(?:[a-z0-9_.-]+/)*(?:ft:)?(?:o[1-9]\\d*(?![a-z0-9])|gpt-[5-9](?:\\.\\d+)?(?![0-9.])|(?:gpt-\\d+(?:\\.\\d+)?(?:-[a-z0-9]+)*-)?(?:codex|deep-research|chat-latest)(?![a-z0-9]))", - "fill_missing_fields": true, + "fill_missing_for_providers": ["openai"], "description": "OpenAI reasoning families by id shape, under any provider namespace and with an optional ft: prefix: the o-series (o1, o3-pro, o4-mini), gpt-5 through gpt-9 majors including dotted minors and suffixed variants (gpt-5.5-cyber, gpt-6-astra), and the codex, deep-research and chat-latest lines when standalone or on a gpt base. gpt-5-search-api is excluded because it is a search-only surface. Every model here is a reasoning model, and the Responses API drops the caller's reasoning param for any mapped OpenAI model whose info lacks supports_reasoning, so an id the registry has not named yet keeps its reasoning settings instead of silently losing them. Rules lose to exact entries. Carries no mode and no pricing, so cost stays on the standard unpriced behavior.", "model_info": { "supports_reasoning": true diff --git a/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py b/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py index eaecbd0b87e..09870faaaf7 100644 --- a/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py +++ b/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py @@ -123,12 +123,13 @@ def test_fill_missing_requires_per_rule_opt_in(restore_generalizations): { "name": "opt-in", "pattern": r"^acme-", - "fill_missing_fields": True, + "fill_missing_for_providers": ["openai"], "model_info": {"supports_vision": True}, }, ] ) - assert match_fill_missing_generalizations("acme-1") == {"supports_vision": True} + assert match_fill_missing_generalizations("acme-1", "openai") == {"supports_vision": True} + assert match_fill_missing_generalizations("acme-1", "azure") is None assert match_capability_generalizations("acme-1") == { "supports_reasoning": True, "supports_vision": True, @@ -137,31 +138,43 @@ def test_fill_missing_requires_per_rule_opt_in(restore_generalizations): restore_generalizations( [{"name": "base", "pattern": r"^acme-", "model_info": {"supports_reasoning": True}}] ) - assert match_fill_missing_generalizations("acme-1") is None + assert match_fill_missing_generalizations("acme-1", "openai") is None restore_generalizations( [ { "name": "mixed", "pattern": r"^acme-", - "fill_missing_fields": True, + "fill_missing_for_providers": ["openai"], "model_info": {"litellm_provider": "openai", "supports_vision": True}, } ] ) - assert match_fill_missing_generalizations("acme-1") == {"supports_vision": True} + assert match_fill_missing_generalizations("acme-1", "openai") == {"supports_vision": True} restore_generalizations( [ { "name": "route", "pattern": r"^acme-", - "fill_missing_fields": True, + "fill_missing_for_providers": ["openai"], "model_info": {"litellm_provider": "openai"}, } ] ) - assert match_fill_missing_generalizations("acme-1") is None + assert match_fill_missing_generalizations("acme-1", "openai") is None + + restore_generalizations( + [ + { + "name": "malformed", + "pattern": r"^acme-", + "fill_missing_for_providers": "openai", + "model_info": {"supports_vision": True}, + } + ] + ) + assert match_fill_missing_generalizations("acme-1", "openai") is None def test_routing_rules_are_excluded_from_capability_results(restore_generalizations): @@ -373,6 +386,12 @@ def test_exact_entries_fill_only_missing_fields(restore_generalizations, monkeyp "litellm_provider": "openai", "mode": "image_generation", }, + "acme-other": { + "input_cost_per_token": 7e-6, + "output_cost_per_token": 8e-6, + "litellm_provider": "openrouter", + "mode": "chat", + }, }, ) restore_generalizations( @@ -380,7 +399,7 @@ def test_exact_entries_fill_only_missing_fields(restore_generalizations, monkeyp { "name": "acme-backfill", "pattern": r"^acme-", - "fill_missing_fields": True, + "fill_missing_for_providers": ["openai"], "model_info": {"supports_reasoning": True, "max_tokens": 5}, } ] @@ -397,6 +416,9 @@ def test_exact_entries_fill_only_missing_fields(restore_generalizations, monkeyp assert bare["input_cost_per_token"] == 3e-6 assert bare["key"] == "acme-bare" + other = litellm.get_model_info("acme-other", custom_llm_provider="openrouter") + assert other.get("supports_reasoning") is None + image = litellm.get_model_info("acme-image", custom_llm_provider="openai") assert image.get("supports_reasoning") is None @@ -727,7 +749,7 @@ def test_shipped_wandb_rule_loses_to_mapped_non_reasoning_entries(shipped_cost_m def test_shipped_wandb_rule_does_not_fill_missing_mapped_entries(shipped_cost_map): - assert match_fill_missing_generalizations("wandb/meta-llama/Llama-3.1-8B-Instruct") is None + assert match_fill_missing_generalizations("wandb/meta-llama/Llama-3.1-8B-Instruct", "wandb") is None def test_shipped_wandb_rule_is_anchored_to_the_wandb_namespace(shipped_cost_map): @@ -854,18 +876,25 @@ def test_shipped_openai_reasoning_rule_loses_to_mapped_entries(shipped_cost_map) [ ("azure/us/o1-2024-12-17", "azure"), ("github_copilot/gpt-5", "github_copilot"), + ("openrouter/openai/o1", "openrouter"), + ("perplexity/openai/gpt-5.4-mini", "perplexity"), ], ) -def test_shipped_openai_reasoning_rule_backfills_mapped_entries(shipped_cost_map, model, provider): +def test_shipped_openai_reasoning_rule_does_not_backfill_other_providers(shipped_cost_map, model, provider): assert model in litellm.model_cost raw_entry = litellm.model_cost[model] assert "supports_reasoning" not in raw_entry model_without_provider = model.removeprefix(f"{provider}/") - assert litellm.supports_reasoning(model=model_without_provider, custom_llm_provider=provider) is True info = litellm.get_model_info(model=model_without_provider, custom_llm_provider=provider) + assert info.get("supports_reasoning") is None assert info["input_cost_per_token"] == raw_entry.get("input_cost_per_token", 0) +def test_shipped_openai_reasoning_rule_matches_only_openai(shipped_cost_map): + assert match_fill_missing_generalizations("gpt-5.4", "openai") == {"supports_reasoning": True} + assert match_fill_missing_generalizations("gpt-5.4", "openrouter") is None + + def test_shipped_openai_reasoning_rule_skips_non_text_modes(shipped_cost_map): model = "gemini/deep-research-pro-preview-12-2025" assert model in litellm.model_cost @@ -877,7 +906,7 @@ def test_shipped_openai_reasoning_rule_skips_non_text_modes(shipped_cost_map): assert info.get("supports_reasoning") is None -def test_shipped_claude_thinking_rules_backfill_without_family_limits(shipped_cost_map): +def test_shipped_claude_thinking_rules_backfill_only_anthropic(shipped_cost_map): model = "perplexity/anthropic/claude-sonnet-4-6" assert model in litellm.model_cost raw_entry = litellm.model_cost[model] @@ -885,6 +914,11 @@ def test_shipped_claude_thinking_rules_backfill_without_family_limits(shipped_co assert "max_input_tokens" not in raw_entry info = litellm.get_model_info(model="anthropic/claude-sonnet-4-6", custom_llm_provider="perplexity") - assert info["supports_adaptive_thinking"] is True - assert info["supports_legacy_thinking"] is True + assert info.get("supports_adaptive_thinking") is None + assert info.get("supports_legacy_thinking") is None assert info.get("max_input_tokens") is None + assert match_fill_missing_generalizations("claude-sonnet-4-6", "anthropic") == { + "supports_adaptive_thinking": True, + "supports_legacy_thinking": True, + } + assert match_fill_missing_generalizations("claude-sonnet-4-6", "perplexity") is None From 71744fb9d19beff93d60cd5d49525e09f3382b8a Mon Sep 17 00:00:00 2001 From: kerry Date: Mon, 14 Sep 2026 17:04:29 +0000 Subject: [PATCH 064/112] style(model_info): format provider backfill call Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/utils.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index 6af66010bd7..aeeec669d24 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5822,9 +5822,7 @@ def _get_model_info_helper( _model_info = None if _model_info is not None and key is not None and _model_info.get("mode", "chat") in _BACKFILL_MODES: - fill_missing: Final = match_fill_missing_generalizations( - key, _model_info.get("litellm_provider", "") - ) + fill_missing: Final = match_fill_missing_generalizations(key, _model_info.get("litellm_provider", "")) if fill_missing is not None: _model_info = { **{k: v for k, v in fill_missing.items() if k not in _model_info}, From 6f086f74b32ce1b2a36790f11902f110224c344c Mon Sep 17 00:00:00 2001 From: kerry Date: Mon, 14 Sep 2026 17:10:54 +0000 Subject: [PATCH 065/112] style(model_info): parameterize fill_missing annotations Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/fallback_generalizations.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/litellm/litellm_core_utils/fallback_generalizations.py b/litellm/litellm_core_utils/fallback_generalizations.py index da75faa4136..7ea4f3589d8 100644 --- a/litellm/litellm_core_utils/fallback_generalizations.py +++ b/litellm/litellm_core_utils/fallback_generalizations.py @@ -183,7 +183,7 @@ class _FallbackGeneralizations: self.rules: list = [] self.routing_rules: tuple = () self.capability_rules: tuple = () - self.fill_missing_rules: tuple = () + self.fill_missing_rules: tuple[_CapabilityRule, ...] = () def set_rules(self, rules: list | None) -> None: installed: Final = rules if isinstance(rules, list) else [] @@ -209,7 +209,7 @@ class _FallbackGeneralizations: return None return {key: value for model_info in matched for key, value in model_info.items()} - def match_fill_missing(self, model: str, provider: str) -> dict | None: + def match_fill_missing(self, model: str, provider: str) -> dict[str, object] | None: if not model or not provider: return None matched = tuple( @@ -219,7 +219,7 @@ class _FallbackGeneralizations: ) if not matched: return None - fill_missing: Final = { + fill_missing: Final[dict[str, object]] = { key: value for model_info in matched for key, value in model_info.items() if key != PROVIDER_KEY } return fill_missing or None @@ -261,7 +261,7 @@ def match_capability_generalizations(model: str) -> dict | None: return _registry.match_capabilities(model) -def match_fill_missing_generalizations(model: str, provider: str) -> dict | None: +def match_fill_missing_generalizations(model: str, provider: str) -> dict[str, object] | None: """Return flagged capability rules matching ``model`` for ``provider``. Later rules override earlier ones on key conflicts. Only rules listing From e4e6b5daeb4a9e7a0acf4236ef63476a402a57eb Mon Sep 17 00:00:00 2001 From: kerry Date: Mon, 14 Sep 2026 17:27:19 +0000 Subject: [PATCH 066/112] fix(lint): allow provider backfill dict annotations Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/fallback_generalizations.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/litellm/litellm_core_utils/fallback_generalizations.py b/litellm/litellm_core_utils/fallback_generalizations.py index 7ea4f3589d8..feea52aa0bc 100644 --- a/litellm/litellm_core_utils/fallback_generalizations.py +++ b/litellm/litellm_core_utils/fallback_generalizations.py @@ -209,7 +209,9 @@ class _FallbackGeneralizations: return None return {key: value for model_info in matched for key, value in model_info.items()} - def match_fill_missing(self, model: str, provider: str) -> dict[str, object] | None: + def match_fill_missing( + self, model: str, provider: str + ) -> dict[str, object] | None: # mutable-ok: preserve the existing dict return contract if not model or not provider: return None matched = tuple( @@ -219,7 +221,7 @@ class _FallbackGeneralizations: ) if not matched: return None - fill_missing: Final[dict[str, object]] = { + fill_missing: Final[dict[str, object]] = { # mutable-ok: preserve the existing dict merge input key: value for model_info in matched for key, value in model_info.items() if key != PROVIDER_KEY } return fill_missing or None @@ -261,7 +263,9 @@ def match_capability_generalizations(model: str) -> dict | None: return _registry.match_capabilities(model) -def match_fill_missing_generalizations(model: str, provider: str) -> dict[str, object] | None: +def match_fill_missing_generalizations( + model: str, provider: str +) -> dict[str, object] | None: # mutable-ok: preserve the existing dict return contract """Return flagged capability rules matching ``model`` for ``provider``. Later rules override earlier ones on key conflicts. Only rules listing From 5781678477196b7d26f0d20682e02ce4b74b1fac Mon Sep 17 00:00:00 2001 From: kerry Date: Mon, 14 Sep 2026 17:33:28 +0000 Subject: [PATCH 067/112] refactor(model_info): drop litellm._logging import from fallback_generalizations Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/fallback_generalizations.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/litellm_core_utils/fallback_generalizations.py b/litellm/litellm_core_utils/fallback_generalizations.py index feea52aa0bc..9b5c369b56e 100644 --- a/litellm/litellm_core_utils/fallback_generalizations.py +++ b/litellm/litellm_core_utils/fallback_generalizations.py @@ -51,12 +51,12 @@ Rules are compiled and classified once, at install time. The match functions are O(number of rules); callers must only invoke them on a cache miss. """ +import logging import re from dataclasses import dataclass from typing import Final -from litellm._logging import verbose_logger - +verbose_logger: Final = logging.getLogger("LiteLLM") NAME_FIELD: Final = "name" PATTERN_FIELD: Final = "pattern" MODEL_INFO_FIELD: Final = "model_info" From e366d3bd302e82ad85c58190f7e23a88b8d56b6b Mon Sep 17 00:00:00 2001 From: kerry Date: Mon, 14 Sep 2026 17:35:24 +0000 Subject: [PATCH 068/112] refactor(model_info): return Mapping from fill_missing matchers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../fallback_generalizations.py | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/litellm/litellm_core_utils/fallback_generalizations.py b/litellm/litellm_core_utils/fallback_generalizations.py index 9b5c369b56e..25ea1bff613 100644 --- a/litellm/litellm_core_utils/fallback_generalizations.py +++ b/litellm/litellm_core_utils/fallback_generalizations.py @@ -53,7 +53,9 @@ O(number of rules); callers must only invoke them on a cache miss. import logging import re +from collections.abc import Mapping from dataclasses import dataclass +from types import MappingProxyType from typing import Final verbose_logger: Final = logging.getLogger("LiteLLM") @@ -209,9 +211,7 @@ class _FallbackGeneralizations: return None return {key: value for model_info in matched for key, value in model_info.items()} - def match_fill_missing( - self, model: str, provider: str - ) -> dict[str, object] | None: # mutable-ok: preserve the existing dict return contract + def match_fill_missing(self, model: str, provider: str) -> Mapping[str, object] | None: if not model or not provider: return None matched = tuple( @@ -221,9 +221,9 @@ class _FallbackGeneralizations: ) if not matched: return None - fill_missing: Final[dict[str, object]] = { # mutable-ok: preserve the existing dict merge input - key: value for model_info in matched for key, value in model_info.items() if key != PROVIDER_KEY - } + fill_missing: Final[Mapping[str, object]] = MappingProxyType( + {key: value for model_info in matched for key, value in model_info.items() if key != PROVIDER_KEY} + ) return fill_missing or None @@ -263,9 +263,7 @@ def match_capability_generalizations(model: str) -> dict | None: return _registry.match_capabilities(model) -def match_fill_missing_generalizations( - model: str, provider: str -) -> dict[str, object] | None: # mutable-ok: preserve the existing dict return contract +def match_fill_missing_generalizations(model: str, provider: str) -> Mapping[str, object] | None: """Return flagged capability rules matching ``model`` for ``provider``. Later rules override earlier ones on key conflicts. Only rules listing From 0f846cecb8566c2fea3e35bb948349fa3c4f7a07 Mon Sep 17 00:00:00 2001 From: kerry Date: Mon, 14 Sep 2026 18:00:15 +0000 Subject: [PATCH 069/112] fix(model_info): single-assign fill_missing_for_providers in rule compiler Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../fallback_generalizations.py | 34 +++++++++++-------- 1 file changed, 19 insertions(+), 15 deletions(-) diff --git a/litellm/litellm_core_utils/fallback_generalizations.py b/litellm/litellm_core_utils/fallback_generalizations.py index 25ea1bff613..5565e557392 100644 --- a/litellm/litellm_core_utils/fallback_generalizations.py +++ b/litellm/litellm_core_utils/fallback_generalizations.py @@ -112,6 +112,22 @@ class _CapabilityRule: _CompiledRule = _RoutingRule | _CapabilityRule +def _parse_fill_missing_for_providers(rule: Mapping[str, object], pattern_label: object) -> frozenset[str] | None: + if FILL_MISSING_FOR_PROVIDERS_FIELD not in rule: + return frozenset() + raw_fill_missing_for_providers: Final = rule.get(FILL_MISSING_FOR_PROVIDERS_FIELD) + if not isinstance(raw_fill_missing_for_providers, (list, tuple)) or not all( + isinstance(provider, str) for provider in raw_fill_missing_for_providers + ): + verbose_logger.warning( + "LiteLLM: skipping malformed fallback generalization rule %s ('%s' must be a list of provider strings).", + rule.get(NAME_FIELD, pattern_label), + FILL_MISSING_FOR_PROVIDERS_FIELD, + ) + return None + return frozenset(raw_fill_missing_for_providers) + + def _compile_rule(rule: object) -> tuple[_CompiledRule, ...]: if not isinstance(rule, dict): return () @@ -134,21 +150,9 @@ def _compile_rule(rule: object) -> tuple[_CompiledRule, ...]: e, ) return () - if FILL_MISSING_FOR_PROVIDERS_FIELD not in rule: - fill_missing_for_providers: Final[frozenset[str]] = frozenset() - else: - raw_fill_missing_for_providers: Final = rule.get(FILL_MISSING_FOR_PROVIDERS_FIELD) - if not isinstance(raw_fill_missing_for_providers, (list, tuple)) or not all( - isinstance(provider, str) for provider in raw_fill_missing_for_providers - ): - verbose_logger.warning( - "LiteLLM: skipping malformed fallback generalization rule %s " - "('%s' must be a list of provider strings).", - rule.get(NAME_FIELD, pattern), - FILL_MISSING_FOR_PROVIDERS_FIELD, - ) - return () - fill_missing_for_providers = frozenset(raw_fill_missing_for_providers) + fill_missing_for_providers: Final = _parse_fill_missing_for_providers(rule, pattern) + if fill_missing_for_providers is None: + return () if PROVIDER_KEY not in model_info: return ( _CapabilityRule( From 4cf20f31ed0067eb97b53b11bbc1df7815a01cd9 Mon Sep 17 00:00:00 2001 From: kerry Date: Mon, 14 Sep 2026 18:39:27 +0000 Subject: [PATCH 070/112] refactor(model_info): drop types import from fallback_generalizations Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/fallback_generalizations.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/litellm/litellm_core_utils/fallback_generalizations.py b/litellm/litellm_core_utils/fallback_generalizations.py index 5565e557392..be432d1fdd4 100644 --- a/litellm/litellm_core_utils/fallback_generalizations.py +++ b/litellm/litellm_core_utils/fallback_generalizations.py @@ -55,7 +55,6 @@ import logging import re from collections.abc import Mapping from dataclasses import dataclass -from types import MappingProxyType from typing import Final verbose_logger: Final = logging.getLogger("LiteLLM") @@ -225,9 +224,9 @@ class _FallbackGeneralizations: ) if not matched: return None - fill_missing: Final[Mapping[str, object]] = MappingProxyType( - {key: value for model_info in matched for key, value in model_info.items() if key != PROVIDER_KEY} - ) + fill_missing: Final[Mapping[str, object]] = { + key: value for model_info in matched for key, value in model_info.items() if key != PROVIDER_KEY + } return fill_missing or None From 250ff03a04ce260ed75d26e28bede2f64091c1f1 Mon Sep 17 00:00:00 2001 From: kerry Date: Mon, 14 Sep 2026 18:53:22 +0000 Subject: [PATCH 071/112] feat(model_info): scope fill_missing rules to azure, bedrock and vertex hosts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../model_prices_and_context_window_backup.json | 8 ++++---- model_prices_and_context_window.json | 8 ++++---- .../test_fallback_generalizations.py | 16 +++++++++------- 3 files changed, 17 insertions(+), 15 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index c9e85c28420..685e63dcc56 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -57717,7 +57717,7 @@ { "name": "claude-adaptive-thinking", "pattern": "claude-[a-z]+-(?:4[-._](?:[6-9]|[1-9]\\d)(?!\\d)|[5-9](?!\\d)(?:[-._]\\d{1,2}(?!\\d))?)", - "fill_missing_for_providers": ["anthropic"], + "fill_missing_for_providers": ["anthropic", "azure_ai", "bedrock", "bedrock_converse", "vertex_ai-anthropic_models"], "description": "Claude at version 4.6 or higher, in any id shape that contains claude--: minors 4.6 through 4.99, any later major-minor, and bare 5+ majors so a new family shaped like claude-fable-5 matches. Requiring the claude- prefix keeps non-Claude names such as team-sonnet-5-1 out. The minor is capped at two digits so an 8-digit date suffix such as claude-opus-4-20250514 is never read as a >= 4.6 minor. Two-digit majors are deliberately not matched so ids like claude-opus-41 (4.1) are not read as major 41. Turns on adaptive thinking for new versions and new families with no code change.", "model_info": { "supports_adaptive_thinking": true @@ -57726,7 +57726,7 @@ { "name": "claude-legacy-thinking", "pattern": "claude-[a-z]+-4[-._]6(?!\\d)", - "fill_missing_for_providers": ["anthropic"], + "fill_missing_for_providers": ["anthropic", "azure_ai", "bedrock", "bedrock_converse", "vertex_ai-anthropic_models"], "description": "Claude at version 4.6 exactly, in any id shape that contains claude--4-6 (dotted and underscored minors included, dated releases such as claude-sonnet-4-6-20260219 too). The 4.6 family is adaptive-thinking yet still accepts legacy thinking.type=enabled with budget_tokens, so the caller's hard budget cap is forwarded verbatim instead of being rewritten to an uncapped output_config.effort. The lookahead keeps two-digit minors such as 4-60 from matching. 4.7+ and 5+ majors reject the legacy shape and stay on the adaptive translation.", "model_info": { "supports_legacy_thinking": true @@ -57743,7 +57743,7 @@ { "name": "claude-mid-conversation-system", "pattern": "claude-[a-z]+-(?:4[-._](?:[89]|[1-9]\\d)(?!\\d)|[5-9](?!\\d)(?:[-._]\\d{1,2}(?!\\d))?)", - "fill_missing_for_providers": ["anthropic"], + "fill_missing_for_providers": ["anthropic", "azure_ai", "bedrock", "bedrock_converse", "vertex_ai-anthropic_models"], "description": "Claude at version 4.8 or higher, in any id shape that contains claude--: minors 4.8 through 4.99, any later major-minor, and bare 5+ majors so a new family like claude-fable-5 matches. Two-digit majors are deliberately not matched so ids like claude-opus-41 (4.1) are not read as major 41. Anthropic introduced mid-conversation system messages with Opus 4.8 and every newer Claude keeps them; 4.7 and below reject the system role inside messages.", "model_info": { "supports_mid_conversation_system": true @@ -57760,7 +57760,7 @@ { "name": "openai-reasoning-family-baseline", "pattern": "^(?!.*search-api)(?:[a-z0-9_.-]+/)*(?:ft:)?(?:o[1-9]\\d*(?![a-z0-9])|gpt-[5-9](?:\\.\\d+)?(?![0-9.])|(?:gpt-\\d+(?:\\.\\d+)?(?:-[a-z0-9]+)*-)?(?:codex|deep-research|chat-latest)(?![a-z0-9]))", - "fill_missing_for_providers": ["openai"], + "fill_missing_for_providers": ["azure", "azure_ai", "openai"], "description": "OpenAI reasoning families by id shape, under any provider namespace and with an optional ft: prefix: the o-series (o1, o3-pro, o4-mini), gpt-5 through gpt-9 majors including dotted minors and suffixed variants (gpt-5.5-cyber, gpt-6-astra), and the codex, deep-research and chat-latest lines when standalone or on a gpt base. gpt-5-search-api is excluded because it is a search-only surface. Every model here is a reasoning model, and the Responses API drops the caller's reasoning param for any mapped OpenAI model whose info lacks supports_reasoning, so an id the registry has not named yet keeps its reasoning settings instead of silently losing them. Rules lose to exact entries. Carries no mode and no pricing, so cost stays on the standard unpriced behavior.", "model_info": { "supports_reasoning": true diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index c9e85c28420..685e63dcc56 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -57717,7 +57717,7 @@ { "name": "claude-adaptive-thinking", "pattern": "claude-[a-z]+-(?:4[-._](?:[6-9]|[1-9]\\d)(?!\\d)|[5-9](?!\\d)(?:[-._]\\d{1,2}(?!\\d))?)", - "fill_missing_for_providers": ["anthropic"], + "fill_missing_for_providers": ["anthropic", "azure_ai", "bedrock", "bedrock_converse", "vertex_ai-anthropic_models"], "description": "Claude at version 4.6 or higher, in any id shape that contains claude--: minors 4.6 through 4.99, any later major-minor, and bare 5+ majors so a new family shaped like claude-fable-5 matches. Requiring the claude- prefix keeps non-Claude names such as team-sonnet-5-1 out. The minor is capped at two digits so an 8-digit date suffix such as claude-opus-4-20250514 is never read as a >= 4.6 minor. Two-digit majors are deliberately not matched so ids like claude-opus-41 (4.1) are not read as major 41. Turns on adaptive thinking for new versions and new families with no code change.", "model_info": { "supports_adaptive_thinking": true @@ -57726,7 +57726,7 @@ { "name": "claude-legacy-thinking", "pattern": "claude-[a-z]+-4[-._]6(?!\\d)", - "fill_missing_for_providers": ["anthropic"], + "fill_missing_for_providers": ["anthropic", "azure_ai", "bedrock", "bedrock_converse", "vertex_ai-anthropic_models"], "description": "Claude at version 4.6 exactly, in any id shape that contains claude--4-6 (dotted and underscored minors included, dated releases such as claude-sonnet-4-6-20260219 too). The 4.6 family is adaptive-thinking yet still accepts legacy thinking.type=enabled with budget_tokens, so the caller's hard budget cap is forwarded verbatim instead of being rewritten to an uncapped output_config.effort. The lookahead keeps two-digit minors such as 4-60 from matching. 4.7+ and 5+ majors reject the legacy shape and stay on the adaptive translation.", "model_info": { "supports_legacy_thinking": true @@ -57743,7 +57743,7 @@ { "name": "claude-mid-conversation-system", "pattern": "claude-[a-z]+-(?:4[-._](?:[89]|[1-9]\\d)(?!\\d)|[5-9](?!\\d)(?:[-._]\\d{1,2}(?!\\d))?)", - "fill_missing_for_providers": ["anthropic"], + "fill_missing_for_providers": ["anthropic", "azure_ai", "bedrock", "bedrock_converse", "vertex_ai-anthropic_models"], "description": "Claude at version 4.8 or higher, in any id shape that contains claude--: minors 4.8 through 4.99, any later major-minor, and bare 5+ majors so a new family like claude-fable-5 matches. Two-digit majors are deliberately not matched so ids like claude-opus-41 (4.1) are not read as major 41. Anthropic introduced mid-conversation system messages with Opus 4.8 and every newer Claude keeps them; 4.7 and below reject the system role inside messages.", "model_info": { "supports_mid_conversation_system": true @@ -57760,7 +57760,7 @@ { "name": "openai-reasoning-family-baseline", "pattern": "^(?!.*search-api)(?:[a-z0-9_.-]+/)*(?:ft:)?(?:o[1-9]\\d*(?![a-z0-9])|gpt-[5-9](?:\\.\\d+)?(?![0-9.])|(?:gpt-\\d+(?:\\.\\d+)?(?:-[a-z0-9]+)*-)?(?:codex|deep-research|chat-latest)(?![a-z0-9]))", - "fill_missing_for_providers": ["openai"], + "fill_missing_for_providers": ["azure", "azure_ai", "openai"], "description": "OpenAI reasoning families by id shape, under any provider namespace and with an optional ft: prefix: the o-series (o1, o3-pro, o4-mini), gpt-5 through gpt-9 majors including dotted minors and suffixed variants (gpt-5.5-cyber, gpt-6-astra), and the codex, deep-research and chat-latest lines when standalone or on a gpt base. gpt-5-search-api is excluded because it is a search-only surface. Every model here is a reasoning model, and the Responses API drops the caller's reasoning param for any mapped OpenAI model whose info lacks supports_reasoning, so an id the registry has not named yet keeps its reasoning settings instead of silently losing them. Rules lose to exact entries. Carries no mode and no pricing, so cost stays on the standard unpriced behavior.", "model_info": { "supports_reasoning": true diff --git a/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py b/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py index 09870faaaf7..b2cc3ebe4c6 100644 --- a/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py +++ b/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py @@ -872,21 +872,23 @@ def test_shipped_openai_reasoning_rule_loses_to_mapped_entries(shipped_cost_map) @pytest.mark.parametrize( - "model,provider", + "model,provider,expected_supports_reasoning", [ - ("azure/us/o1-2024-12-17", "azure"), - ("github_copilot/gpt-5", "github_copilot"), - ("openrouter/openai/o1", "openrouter"), - ("perplexity/openai/gpt-5.4-mini", "perplexity"), + ("azure/us/o1-2024-12-17", "azure", True), + ("github_copilot/gpt-5", "github_copilot", None), + ("openrouter/openai/o1", "openrouter", None), + ("perplexity/openai/gpt-5.4-mini", "perplexity", None), ], ) -def test_shipped_openai_reasoning_rule_does_not_backfill_other_providers(shipped_cost_map, model, provider): +def test_shipped_openai_reasoning_rule_backfills_only_approved_providers( + shipped_cost_map, model, provider, expected_supports_reasoning +): assert model in litellm.model_cost raw_entry = litellm.model_cost[model] assert "supports_reasoning" not in raw_entry model_without_provider = model.removeprefix(f"{provider}/") info = litellm.get_model_info(model=model_without_provider, custom_llm_provider=provider) - assert info.get("supports_reasoning") is None + assert info.get("supports_reasoning") is expected_supports_reasoning assert info["input_cost_per_token"] == raw_entry.get("input_cost_per_token", 0) From ce0301c23f0658e446cdc446602ed6cada162b94 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 19:00:49 +0000 Subject: [PATCH 072/112] refactor(proxy): drop redundant docstrings from upload allowlist helpers and tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../openai_files_endpoints/general_upload_validation.py | 5 ----- .../proxy/openai_files_endpoint/test_files_endpoint.py | 1 - .../openai_files_endpoint/test_general_upload_validation.py | 1 - 3 files changed, 7 deletions(-) diff --git a/litellm/proxy/openai_files_endpoints/general_upload_validation.py b/litellm/proxy/openai_files_endpoints/general_upload_validation.py index 22f34b8bd87..e9b6c319fe5 100644 --- a/litellm/proxy/openai_files_endpoints/general_upload_validation.py +++ b/litellm/proxy/openai_files_endpoints/general_upload_validation.py @@ -32,10 +32,6 @@ def coerce_optional_int_setting(raw: object) -> int | None: def coerce_optional_str_list_setting(raw: object) -> tuple[str, ...] | None: - """A general_settings value declared as an optional list of strings, e.g. allowed_file_extensions. - - None (unset) and [] (set to nothing) are different answers for an allowlist, so both survive. - """ if raw is None: return None if not isinstance(raw, list) or not all(isinstance(item, str) for item in raw): @@ -104,7 +100,6 @@ def check_allowed_extension( filename: str | None, allowed_extensions: tuple[str, ...] | None, ) -> UploadedFileExtensionNotAllowed | None: - """None means the allowlist is not configured; an empty tuple means nothing is allowed.""" if allowed_extensions is None: return None extension: Final = _normalized_extension(filename) diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index bd3aafd7865..5d8222162a2 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -4776,7 +4776,6 @@ def test_create_file_empty_allowlist_rejects_every_upload(monkeypatch, llm_route def test_create_file_allowlist_runs_before_blocklist(monkeypatch, llm_router: Router): - """An extension in both lists is refused by the allowlist message, and the blocklist still holds on its own.""" import litellm.proxy.proxy_server as ps forwarded_calls = _setup_batch_upload_endpoint(monkeypatch, llm_router) diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_general_upload_validation.py b/tests/test_litellm/proxy/openai_files_endpoint/test_general_upload_validation.py index f5b558f3334..b742e1aa9b6 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_general_upload_validation.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_general_upload_validation.py @@ -104,7 +104,6 @@ def test_allowed_extension_match_is_case_insensitive_for_configured_value(): @pytest.mark.parametrize("filename", ["README", "", None, "../../"]) def test_no_extension_rejected_when_allowlist_set(filename): - """The allowlist grants by extension, so a name that yields none has nothing to be granted for.""" assert check_allowed_extension(filename, (".jsonl",)) == UploadedFileExtensionNotAllowed(extension="") From 209fc7afb0800c303b8e5bdcb7fbd46d59121ad9 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 14 Sep 2026 12:03:54 -0700 Subject: [PATCH 073/112] fix(harness): trace OCR Python and native dispatch paths --- .../strategies/trace_parity/models.py | 1 + .../strategies/trace_parity/sdk/execution.py | 22 ++- .../strategies/trace_parity/sdk/ocr/case.py | 128 ++++++++++-------- .../sdk/test_core_scenario_matrix.py | 6 + .../strategies/trace_parity/test_runner.py | 33 ++++- 5 files changed, 132 insertions(+), 58 deletions(-) diff --git a/tests/rust-python-harness/strategies/trace_parity/models.py b/tests/rust-python-harness/strategies/trace_parity/models.py index 078af2b316d..be7375dc604 100644 --- a/tests/rust-python-harness/strategies/trace_parity/models.py +++ b/tests/rust-python-harness/strategies/trace_parity/models.py @@ -64,6 +64,7 @@ class TraceScenario: fixture: Callable[[Engine, str], RouteFixture] mappings: tuple[TraceMapping, ...] asynchronous: bool + python_rust_enabled: bool = False @dataclass(frozen=True, slots=True) diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py b/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py index c249d09fa73..f8a674451b8 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py @@ -1,10 +1,13 @@ from __future__ import annotations import asyncio +import os from collections.abc import AsyncIterable, Awaitable, Iterable +from contextlib import nullcontext from dataclasses import dataclass from pathlib import Path from typing import Final, Protocol, cast +from unittest.mock import patch from ....shared.parity.replay import replay_server from ....shared.reporting.models import Surface @@ -105,7 +108,7 @@ def _collect( def collect_trace( - spec: RouteSpec, engine: Engine, *, asynchronous: bool + spec: RouteSpec, engine: Engine, *, asynchronous: bool, python_rust_enabled: bool = False ) -> tuple[FunctionTraceEvent, ...] | TraceExecutionFailure: function: Final = _entrypoint(spec, engine, asynchronous=asynchronous) if isinstance(function, TraceExecutionFailure): @@ -126,7 +129,13 @@ def collect_trace( expected_failure=base_fixture.expected_failure, consume_stream=base_fixture.consume_stream, ) - collected: Final = _collect(function, fixture, engine, asynchronous=asynchronous) + environment: Final = ( + patch.dict(os.environ, {"LITELLM_RUST": "1" if python_rust_enabled else "0"}) + if engine == "python" + else nullcontext() + ) + with environment: + collected: Final = _collect(function, fixture, engine, asynchronous=asynchronous) provider.take_requests(len(fixture.provider_responses)) except Exception as error: return TraceExecutionFailure(engine, f"{type(error).__name__}: {error}") @@ -159,7 +168,14 @@ def execute_trace( fixture=scenario.fixture, ) python_trace: Final = ( - collect_trace(scenario_route, "python", asynchronous=scenario.asynchronous) if engine != "rust" else () + collect_trace( + scenario_route, + "python", + asynchronous=scenario.asynchronous, + python_rust_enabled=scenario.python_rust_enabled, + ) + if engine != "rust" + else () ) rust_trace: Final = ( collect_trace(scenario_route, "rust", asynchronous=scenario.asynchronous) if engine != "python" else () diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py b/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py index effd1a0b4f6..b2c0e4c7acc 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py @@ -1,7 +1,7 @@ from __future__ import annotations import json -from typing import Final, cast +from typing import Final from .....shared.parity.recorded_http import HttpHeader, RecordedHttpResponse from .....shared.tracing.steps import Engine, mapping @@ -53,6 +53,23 @@ ASYNC_MAPPINGS: Final = ( ), ) +PUBLIC_RUST_DISPATCH_MAPPINGS: Final = ( + mapping(span="public_sdk_entrypoint", python_frame=r"ocr/main\.py:\d+ a?ocr$"), + mapping(span="public_request", python_frame=r"ocr/main\.py:\d+ _public_request$"), + mapping(span="bind_request", python_frame=r"ocr/main\.py:\d+ _bind_request$"), + mapping(span="rust_ocr_enabled", python_frame=r"rust_bridge/configuration\.py:\d+ rust_ocr_enabled$"), + mapping(span="select_native_ocr", python_frame=r"rust_bridge/ocr_lifecycle\.py:\d+ select$"), + mapping(span="load_native_bridge", python_frame=r"rust_bridge/bindings\.py:\d+ NativeBinding\.load$"), + mapping(span="native_call_setup", python_frame=r"rust_bridge/lifecycle\.py:\d+ setup$"), + mapping(span="native_response", python_frame=r"rust_bridge/ocr\.py:\d+ _response$"), + mapping(span="native_call_finalize", python_frame=r"rust_bridge/lifecycle\.py:\d+ finalize$"), + mapping( + span="native_success_bookkeeping", + python_frame=r"rust_bridge/lifecycle\.py:\d+ success_bookkeeping$", + ), + *(mapping(rust_span=item.rust) for item in SYNC_MAPPINGS if item.rust is not None), +) + CALLBACK_SUCCESS_SYNC_MAPPINGS: Final = (*SYNC_MAPPINGS, SUCCESS_CALLBACK_SYNC_MAPPING) CALLBACK_SUCCESS_ASYNC_MAPPINGS: Final = (*ASYNC_MAPPINGS, SUCCESS_CALLBACK_ASYNC_MAPPING) CALLBACK_FAILURE_SYNC_MAPPINGS: Final = ( @@ -164,23 +181,6 @@ def _azure_fixture(engine: Engine, _base_url: str) -> RouteFixture: ) -def _vertex_fixture(engine: Engine, _base_url: str) -> RouteFixture: - fixture: Final = _fixture( - engine, - "vertex_ai/mistral-ocr-maas", - {"type": "image_url", "image_url": "data:image/png;base64,aGVsbG8="}, - ) - vertex: Final = {"vertex_project": "trace-project", "vertex_location": "us-central1"} - optional_params: Final = cast(dict[str, object], fixture.kwargs.get("optional_params", {})) - return RouteFixture( - kwargs={ - **fixture.kwargs, - **({"optional_params": {**optional_params, **vertex}} if engine == "rust" else vertex), - }, - provider_responses=fixture.provider_responses, - ) - - def _vertex_deepseek_fixture(engine: Engine, _base_url: str) -> RouteFixture: vertex: Final = {"vertex_project": "trace-project", "vertex_location": "us-central1"} return RouteFixture( @@ -204,6 +204,28 @@ def _vertex_deepseek_fixture(engine: Engine, _base_url: str) -> RouteFixture: ) +def _cohere_fixture(engine: Engine, _base_url: str) -> RouteFixture: + return RouteFixture( + kwargs={ + "model": "cohere/parse-v5.0", + "document": {"type": "image_url", "image_url": "data:image/png;base64,aGVsbG8="}, + **({"optional_params": {"output_format": "blocks"}} if engine == "rust" else {"output_format": "blocks"}), + }, + provider_responses=( + RecordedHttpResponse.from_bytes( + 200, + (HttpHeader(name="content-type", value="application/json"),), + json.dumps( + { + "pages": [{"index": 0, "blocks": [{"type": "text", "text": {"content": "hello"}}]}], + "meta": {"billed_units": {"pages": 1}}, + } + ).encode(), + ), + ), + ) + + def _azure_document_intelligence_fixture(engine: Engine, base_url: str) -> RouteFixture: completed: Final = json.dumps( { @@ -249,30 +271,6 @@ def _azure_document_intelligence_fixture(engine: Engine, base_url: str) -> Route ) -VERTEX_COMMON_MAPPINGS: Final = ( - *COMMON_MAPPINGS[:7], - mapping( - rust_span="transform_ocr_request", - python_frame=( - r"VertexAIOCRConfig\.(?:async_)?transform_ocr_request$" - r"|MistralOCRConfig\.transform_ocr_request$" - ), - ), - COMMON_MAPPINGS[-1], -) -VERTEX_SYNC_MAPPINGS: Final = ( - *VERTEX_COMMON_MAPPINGS, - mapping(rust_span="execute_ocr_provider_call", python_frame=r"BaseLLMHTTPHandler\.ocr$"), - mapping(span="python_transform_ocr_response_wrapper", python_frame=r"BaseLLMHTTPHandler\._transform_ocr_response$"), - mapping(rust_span="transform_ocr_response", python_frame=r"MistralOCRConfig\.transform_ocr_response$"), -) -VERTEX_ASYNC_MAPPINGS: Final = ( - *VERTEX_COMMON_MAPPINGS, - mapping(span="python_ocr_wrapper", python_frame=r"BaseLLMHTTPHandler\.ocr$"), - mapping(rust_span="execute_ocr_provider_call", python_frame=r"BaseLLMHTTPHandler\.async_ocr$"), - mapping(rust_span="transform_ocr_response", python_frame=r"MistralOCRConfig\.transform_ocr_response$"), -) - DEEPSEEK_COMMON_MAPPINGS: Final = ( mapping(rust_span="ocr", python_frame=r"ocr/main\.py:\d+ a?ocr$"), mapping(rust_span="prepare_ocr_call", python_frame=r"ocr/main\.py:\d+ _prepare_ocr_request$"), @@ -352,6 +350,20 @@ DOCUMENT_INTELLIGENCE_ASYNC_MAPPINGS: Final = ( mapping(span="python_poll_http_request", python_frame=r"AsyncHTTPHandler\.get$"), ) +COHERE_COMMON_MAPPINGS: Final = ( + *COMMON_MAPPINGS[:7], + mapping( + rust_span="transform_ocr_request", + python_frame=r"CohereParseConfig\.(?:async_)?transform_ocr_request$", + ), + COMMON_MAPPINGS[-1], + mapping(rust_span="transform_ocr_response", python_frame=r"CohereParseConfig\.transform_ocr_response$"), +) +COHERE_ASYNC_MAPPINGS: Final = ( + *COHERE_COMMON_MAPPINGS, + mapping(span="python_ocr_wrapper", python_frame=r"BaseLLMHTTPHandler\.ocr$"), + mapping(rust_span="execute_ocr_provider_call", python_frame=r"BaseLLMHTTPHandler\.async_ocr$"), +) SPEC: Final = RouteSpec("ocr", ("ocr", "aocr"), ("ocr", "aocr"), _mistral_fixture) TRACE_SUITE: Final = TraceSuite( @@ -417,18 +429,6 @@ TRACE_SUITE: Final = TraceSuite( mappings=(*DOCUMENT_INTELLIGENCE_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), asynchronous=True, ), - TraceScenario( - name="sync-vertex-ai", - fixture=_vertex_fixture, - mappings=(*VERTEX_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), - asynchronous=False, - ), - TraceScenario( - name="async-vertex-ai", - fixture=_vertex_fixture, - mappings=(*VERTEX_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), - asynchronous=True, - ), TraceScenario( name="sync-vertex-deepseek", fixture=_vertex_deepseek_fixture, @@ -441,5 +441,25 @@ TRACE_SUITE: Final = TraceSuite( mappings=(*DEEPSEEK_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), asynchronous=True, ), + TraceScenario( + name="async-cohere", + fixture=_cohere_fixture, + mappings=(*COHERE_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), + asynchronous=True, + ), + TraceScenario( + name="sync-public-rust-dispatch", + fixture=_mistral_fixture, + mappings=PUBLIC_RUST_DISPATCH_MAPPINGS, + asynchronous=False, + python_rust_enabled=True, + ), + TraceScenario( + name="async-public-rust-dispatch", + fixture=_mistral_fixture, + mappings=PUBLIC_RUST_DISPATCH_MAPPINGS, + asynchronous=True, + python_rust_enabled=True, + ), ), ) diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/test_core_scenario_matrix.py b/tests/rust-python-harness/strategies/trace_parity/sdk/test_core_scenario_matrix.py index d0921398795..d0dbd281a97 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/test_core_scenario_matrix.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/test_core_scenario_matrix.py @@ -16,6 +16,7 @@ def _suite(module: str) -> TraceSuite: def test_core_sdk_scenario_matrix_keeps_distinct_migration_paths() -> None: chat: Final = _suite("tests.rust-python-harness.strategies.trace_parity.sdk.chat_completions.case") messages: Final = _suite("tests.rust-python-harness.strategies.trace_parity.sdk.messages.case") + ocr: Final = _suite("tests.rust-python-harness.strategies.trace_parity.sdk.ocr.case") responses: Final = _suite("tests.rust-python-harness.strategies.trace_parity.sdk.responses.case") assert {(scenario.name, scenario.asynchronous) for scenario in chat.scenarios} >= { @@ -38,6 +39,11 @@ def test_core_sdk_scenario_matrix_keeps_distinct_migration_paths() -> None: ("async-bedrock-invalid-thinking-retry", True), ("sync-unsupported", False), } + assert {(scenario.name, scenario.asynchronous) for scenario in ocr.scenarios} >= { + ("async-cohere", True), + ("sync-public-rust-dispatch", False), + ("async-public-rust-dispatch", True), + } assert {(scenario.name, scenario.asynchronous) for scenario in responses.scenarios} >= { ("sync-openai", False), ("async-openai", True), diff --git a/tests/rust-python-harness/strategies/trace_parity/test_runner.py b/tests/rust-python-harness/strategies/trace_parity/test_runner.py index 4992854e21d..9dfd0d51c89 100644 --- a/tests/rust-python-harness/strategies/trace_parity/test_runner.py +++ b/tests/rust-python-harness/strategies/trace_parity/test_runner.py @@ -1,7 +1,9 @@ from __future__ import annotations import importlib +import os from pathlib import Path +from types import SimpleNamespace from typing import Final, cast import pytest @@ -10,11 +12,12 @@ import litellm from ...shared.reporting.models import Coverage, HarnessCase, HarnessRun, RunStatus, SdkFunction, Surface from ...shared.reporting.strategy import ModuleCaseSpec +from ...shared.tracing.profiler import FunctionTraceEvent from ...shared.tracing.steps import Engine, PipelineStep from .models import GatewayRouteSpec, RouteFixture, RouteSpec, TraceScenario, TraceSuite from .reporting import TraceArtifact from .runner import run_trace_cases, run_trace_scenario, runner_selection, scenario_nodeids, validate_trace_suite -from .sdk.execution import execute_trace +from .sdk.execution import SdkCall, collect_trace, execute_trace def _fixture(_engine: Engine, _base_url: str) -> RouteFixture: @@ -77,6 +80,34 @@ def test_python_engine_skips_native_bridge(monkeypatch: pytest.MonkeyPatch, tmp_ assert selected == [(frozenset({"mistral"}), "python")] +def test_python_trace_controls_native_ocr_dispatch(monkeypatch: pytest.MonkeyPatch) -> None: + execution: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.sdk.execution") + route: Final = RouteSpec("ocr", ("ocr", "aocr"), ("ocr", "aocr"), _fixture) + observed: list[str | None] = [] + + def collect( + _function: SdkCall, + _fixture: RouteFixture, + _engine: Engine, + *, + asynchronous: bool, + ) -> SimpleNamespace: + observed.append(os.environ.get("LITELLM_RUST")) + return SimpleNamespace( + events=(FunctionTraceEvent(0, None, "aocr" if asynchronous else "ocr"),), + error=None, + ) + + monkeypatch.setenv("LITELLM_RUST", "1") + monkeypatch.setattr(execution, "_collect", collect) + + collect_trace(route, "python", asynchronous=False) + collect_trace(route, "python", asynchronous=True, python_rust_enabled=True) + + assert observed == ["0", "1"] + assert os.environ["LITELLM_RUST"] == "1" + + def test_expected_provider_failure_omits_feedback_banner( capsys: pytest.CaptureFixture[str], monkeypatch: pytest.MonkeyPatch ) -> None: From ba171d9eb4a96eead01626ac0c80cb17651c5555 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 19:06:04 +0000 Subject: [PATCH 074/112] fix(proxy): hide default credentials login hint when UI_PASSWORD is set Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/auth_utils.py | 12 ++++++ .../ui_discovery_endpoints.py | 7 +--- litellm/proxy/management_endpoints/ui_sso.py | 6 +-- litellm/proxy/proxy_server.py | 6 +-- .../test_ui_discovery_endpoints.py | 40 +++++++++++++++++++ .../proxy/management_endpoints/test_ui_sso.py | 17 ++++++++ .../proxy_server/test_routes_login_sso.py | 12 ++++++ 7 files changed, 87 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index be65c3b39ec..0f4160d840b 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -1446,6 +1446,18 @@ def has_user_setup_sso() -> bool: ) +def should_hide_default_credentials_hint(general_settings: Mapping[str, object]) -> bool: + """ + Whether login pages hide the "admin / MASTER_KEY" hint: explicit opt-in, or a + non-empty UI_PASSWORD, which makes that hint wrong. UI_USERNAME alone keeps it. + """ + return ( + os.getenv("LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT", "false").lower() == "true" + or general_settings.get("hide_default_credentials_hint", False) is True + or bool(os.getenv("UI_PASSWORD")) + ) + + def _is_google_ready() -> bool: return bool(os.getenv("GOOGLE_CLIENT_ID")) and bool(os.getenv("GOOGLE_CLIENT_SECRET")) diff --git a/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py b/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py index c2053693f2e..07fd04a74fb 100644 --- a/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py +++ b/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py @@ -14,7 +14,7 @@ router: Final = APIRouter() @router.get("/.well-known/litellm-ui-config", response_model=UiDiscoveryEndpoints) @router.get("/litellm/.well-known/litellm-ui-config", response_model=UiDiscoveryEndpoints) # if mounted at root path async def get_ui_config(): - from litellm.proxy.auth.auth_utils import has_user_setup_sso + from litellm.proxy.auth.auth_utils import has_user_setup_sso, should_hide_default_credentials_hint from litellm.proxy.proxy_server import general_settings from litellm.proxy.utils import get_proxy_base_url, get_server_root_path @@ -23,10 +23,7 @@ async def get_ui_config(): or general_settings.get("auto_redirect_ui_login_to_sso", False) is True ) admin_ui_disabled: Final = os.getenv("DISABLE_ADMIN_UI", "false").lower() == "true" - hide_default_credentials_hint: Final = bool( - os.getenv("LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT", "false").lower() == "true" - or general_settings.get("hide_default_credentials_hint", False) is True - ) + hide_default_credentials_hint: Final = should_hide_default_credentials_hint(general_settings) sso_configured: Final = has_user_setup_sso() diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 1ba90725eff..ab916e8df1b 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -92,6 +92,7 @@ from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken, get_user_obje from litellm.proxy.auth.auth_utils import ( _get_request_ip_address, has_user_setup_sso, + should_hide_default_credentials_hint, ) from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.ip_address_utils import IPAddressUtils @@ -1110,10 +1111,7 @@ async def google_login( from fastapi.responses import HTMLResponse - hide_default_credentials_hint: Final = ( - os.getenv("LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT", "false").lower() == "true" - or general_settings.get("hide_default_credentials_hint", False) is True - ) + hide_default_credentials_hint: Final = should_hide_default_credentials_hint(general_settings) form_response: Final = HTMLResponse( content=build_ui_login_form( show_deprecation_banner=True, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d9f8e04ebda..5b3d6365df4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -318,6 +318,7 @@ from litellm.proxy.auth.auth_utils import ( check_response_size_is_safe, is_request_body_safe, log_once_if_budget_reservation_disabled, + should_hide_default_credentials_hint, warn_once_if_custom_auth_skips_common_checks, ) from litellm.proxy.auth.fallback_model_access import router_fallback_access_check @@ -15816,10 +15817,7 @@ async def fallback_login(request: Request): from fastapi.responses import HTMLResponse - hide_default_credentials_hint: Final = ( - os.getenv("LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT", "false").lower() == "true" - or general_settings.get("hide_default_credentials_hint", False) is True - ) + hide_default_credentials_hint: Final = should_hide_default_credentials_hint(general_settings) return HTMLResponse( content=build_ui_login_form( show_deprecation_banner=False, diff --git a/tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py b/tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py index c3f7b0100d8..5e3d9c0e43b 100644 --- a/tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py +++ b/tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py @@ -340,6 +340,7 @@ def test_ui_discovery_endpoints_hide_default_credentials_hint_default_false(): patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}, clear=False), ): os.environ.pop("LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT", None) + os.environ.pop("UI_PASSWORD", None) response = client.get("/.well-known/litellm-ui-config") @@ -348,6 +349,45 @@ def test_ui_discovery_endpoints_hide_default_credentials_hint_default_false(): assert data["hide_default_credentials_hint"] is False +def test_ui_discovery_endpoints_hide_default_credentials_hint_when_ui_password_set(): + """A non-empty UI_PASSWORD hides the card: 'admin / MASTER_KEY' is no longer the login.""" + app = FastAPI() + app.include_router(router) + client = TestClient(app) + + with patch.dict(os.environ, {"UI_PASSWORD": "s3cret-pass", "DISABLE_ADMIN_UI": "false"}, clear=False): + os.environ.pop("LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT", None) + + response = client.get("/.well-known/litellm-ui-config") + + assert response.status_code == 200 + assert response.json()["hide_default_credentials_hint"] is True + + +@pytest.mark.parametrize( + "env_overrides", + [ + pytest.param({"UI_USERNAME": "opsadmin"}, id="username_only_keeps_master_key_password"), + pytest.param({"UI_PASSWORD": ""}, id="empty_password_is_not_set"), + ], +) +def test_ui_discovery_endpoints_keeps_default_credentials_hint_without_real_ui_password(env_overrides): + """Only a non-empty UI_PASSWORD counts as custom credentials; the hint stays accurate otherwise.""" + app = FastAPI() + app.include_router(router) + client = TestClient(app) + + with patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false", **env_overrides}, clear=False): + os.environ.pop("LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT", None) + if "UI_PASSWORD" not in env_overrides: + os.environ.pop("UI_PASSWORD", None) + + response = client.get("/.well-known/litellm-ui-config") + + assert response.status_code == 200 + assert response.json()["hide_default_credentials_hint"] is False + + def test_ui_discovery_endpoints_hide_default_credentials_hint_via_env_var(): """LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT=true hides the login-page credentials card.""" app = FastAPI() diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 2050e65d2a1..2e3c2ceffa3 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -8334,6 +8334,7 @@ async def _render_legacy_login_page(env_overrides, general_settings): "GOOGLE_CLIENT_ID", "GENERIC_CLIENT_ID", "LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT", + "UI_PASSWORD", ): os.environ.pop(var, None) os.environ.update(env_overrides) @@ -8386,6 +8387,22 @@ async def test_legacy_login_page_hides_credentials_hint_via_general_settings(): assert "MASTER_KEY" not in body +@pytest.mark.asyncio +async def test_legacy_login_page_hides_credentials_hint_when_ui_password_set(): + """Regression: the legacy page shares the rule with the discovery endpoint, so a non-empty + UI_PASSWORD hides the now-inaccurate 'admin / MASTER_KEY' hint here too.""" + response = await _render_legacy_login_page( + env_overrides={"UI_PASSWORD": "s3cret-pass"}, + general_settings={}, + ) + + body = response.body.decode() + assert response.status_code == 200 + assert "Default Credentials" not in body + assert "MASTER_KEY" not in body + assert 'name="username"' in body + + @pytest.mark.asyncio async def test_saml_callback_blocked_when_admin_ui_disabled(): """An IdP-initiated assertion must not mint a UI session when the admin UI is diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_login_sso.py b/tests/test_litellm/proxy/proxy_server/test_routes_login_sso.py index 45460dcecf1..1b4f9bc1cfe 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_login_sso.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_login_sso.py @@ -97,6 +97,7 @@ def test_fallback_login_returns_html_form_with_ui_username_set(client, monkeypat def test_fallback_login_shows_credentials_hint_by_default(client, monkeypatch): """Control: without the flag, /fallback/login still renders the hint.""" monkeypatch.delenv("UI_USERNAME", raising=False) + monkeypatch.delenv("UI_PASSWORD", raising=False) monkeypatch.delenv("LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT", raising=False) response = client.get("/fallback/login") assert response.status_code == 200 @@ -104,6 +105,17 @@ def test_fallback_login_shows_credentials_hint_by_default(client, monkeypatch): assert "MASTER_KEY" in response.text +def test_fallback_login_hides_credentials_hint_when_ui_password_set(client, monkeypatch): + """Regression: a non-empty UI_PASSWORD means 'admin / MASTER_KEY' is wrong, so the hint must go.""" + monkeypatch.setenv("UI_PASSWORD", "s3cret-pass") + monkeypatch.delenv("LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT", raising=False) + response = client.get("/fallback/login") + assert response.status_code == 200 + assert "Default Credentials" not in response.text + assert "MASTER_KEY" not in response.text + assert 'name="username"' in response.text + + def test_fallback_login_hides_credentials_hint_via_env_flag(client, monkeypatch): """Pin: LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT removes the hint on /fallback/login.""" monkeypatch.delenv("UI_USERNAME", raising=False) From b00bb3556327ad576bea306bf93440dee27227b0 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 19:06:29 +0000 Subject: [PATCH 075/112] fix(auth): drop the membership write-epoch and background Redis replicate, load membership lazily Move the cache-miss marker out of litellm.constants into auth_checks (CodeQL cyclic import) and stop logging user_id/team_id in the lookup failure (CodeQL log injection). Write the membership row through DualCache synchronously again instead of a background Redis task guarded by a bounded write-epoch map: the epoch was sampled after the Prisma read, so an invalidate that raced the read could be cached as current, and eviction of the epoch entry could let an old Redis write land. The synchronous write keeps invalidate_team_member_spend_state authoritative. Lookup failures return None again (fail-open like main) instead of 503, and the load is skipped on routes that neither resolve a model nor run budget checks. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 7 - litellm/proxy/auth/auth_checks.py | 179 +++----------- .../proxy/auth/test_auth_checks.py | 226 ++++++------------ 3 files changed, 106 insertions(+), 306 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 5c7a02d0743..5751e6e46af 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -2039,10 +2039,3 @@ BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY: Final = "batch_enqueued_token_limit" # Shared read-only empty mapping, for defaulting optional Mapping parameters without # constructing a fresh mutable dict at each call site. EMPTY_MAPPING: Final = MappingProxyType({}) - - -class TeamMembershipCacheMiss: - __slots__ = () - - -TEAM_MEMBERSHIP_CACHE_MISS: Final = TeamMembershipCacheMiss() diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 89b061bdefe..a8f9cc77577 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -35,8 +35,6 @@ from litellm.constants import ( MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE, REGISTRY_ERROR_NEGATIVE_CACHE_TTL, TAG_REGISTRY_MAX_SIZE, - TEAM_MEMBERSHIP_CACHE_MISS, - TeamMembershipCacheMiss, ) from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider @@ -331,27 +329,19 @@ db_cache_expiry: Final = DEFAULT_IN_MEMORY_TTL # refresh every 5s _TEAM_MEMBERSHIP_INFLIGHT_MAX: Final = 10000 _team_membership_inflight: Final = LimitedSizeOrderedDict(max_size=_TEAM_MEMBERSHIP_INFLIGHT_MAX) -_team_membership_write_epoch: Final = LimitedSizeOrderedDict(max_size=_TEAM_MEMBERSHIP_INFLIGHT_MAX) + + +class _TeamMembershipCacheMiss: + __slots__ = () + + +_TEAM_MEMBERSHIP_CACHE_MISS: Final = _TeamMembershipCacheMiss() all_routes: Final = LiteLLMRoutes.openai_routes.value + LiteLLMRoutes.management_routes.value -def _membership_write_epoch(key: str) -> int: - cached: Final[object] = _team_membership_write_epoch.get(key, 0) - return cached if isinstance(cached, int) else 0 - - -def _bump_membership_write_epoch(key: str) -> None: - _team_membership_write_epoch[key] = _membership_write_epoch(key) + 1 - - def _membership_from_shared_load(result: object) -> LiteLLM_TeamMembership | None: - if result is None or isinstance(result, LiteLLM_TeamMembership): - return result - raise HTTPException( - status_code=status.HTTP_503_SERVICE_UNAVAILABLE, - detail="Failed to load team membership", - ) + return result if isinstance(result, LiteLLM_TeamMembership) else None def _log_budget_lookup_failure(entity: str, error: Exception) -> None: @@ -897,21 +887,6 @@ async def common_checks( """ from litellm.proxy.proxy_server import prisma_client, user_api_key_cache - # One membership read for model-access, access-group attribution, and - # member-budget. Each used to call get_team_membership independently; - # DualCache Redis SET/GET on every miss made those look like two Postgres spans. - loaded_team_membership: LiteLLM_TeamMembership | None = None - team_membership_loaded = False - if team_object is not None and valid_token is not None and valid_token.user_id is not None: - loaded_team_membership = await get_team_membership( - user_id=valid_token.user_id, - team_id=team_object.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - team_membership_loaded = True - _model: Final[str | list[str] | None] = get_model_from_request( request_data=request_body, route=route, @@ -926,6 +901,22 @@ async def common_checks( and (route in MODEL_DISCOVERY_ROUTES or not RouteChecks.is_llm_api_route(route=route)) ) + membership_user_id: Final = ( + valid_token.user_id if valid_token is not None and (bool(_model) or not skip_all_budget_checks) else None + ) + team_membership_loaded: Final = team_object is not None and membership_user_id is not None + loaded_team_membership: Final = ( + await get_team_membership( + user_id=membership_user_id, + team_id=team_object.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + if team_object is not None and membership_user_id is not None + else None + ) + unpriced_models: Final = ( _unpriced_models_in_request(model=_model, llm_router=llm_router) if litellm.block_requests_for_models_without_pricing and RouteChecks.is_llm_api_route(route=route) @@ -2188,90 +2179,13 @@ async def get_tag_object( def _membership_from_cached_payload( cached: object, -) -> LiteLLM_TeamMembership | None | TeamMembershipCacheMiss: +) -> LiteLLM_TeamMembership | None | _TeamMembershipCacheMiss: if cached is None: - return TEAM_MEMBERSHIP_CACHE_MISS + return _TEAM_MEMBERSHIP_CACHE_MISS if cached == NO_TEAM_MEMBERSHIP_SENTINEL: return None cached_membership: Final = CacheCodec.deserialize(cached, model_type=LiteLLM_TeamMembership) - return cached_membership if cached_membership is not None else TEAM_MEMBERSHIP_CACHE_MISS - - -async def _set_team_membership_l1( - user_api_key_cache: UserApiKeyCache, - key: str, - value: object, - *, - model_type: type[LiteLLM_TeamMembership] | None, - ttl: float | None, -) -> None: - match (model_type is not None, ttl is not None): - case (False, False): - await user_api_key_cache.async_set_cache(key=key, value=value, local_only=True) - case (False, True): - await user_api_key_cache.async_set_cache(key=key, value=value, local_only=True, ttl=ttl) - case (True, False): - await user_api_key_cache.async_set_cache(key=key, value=value, local_only=True, model_type=model_type) - case (True, True): - await user_api_key_cache.async_set_cache( - key=key, value=value, local_only=True, model_type=model_type, ttl=ttl - ) - - -async def _replicate_team_membership_to_redis( - user_api_key_cache: UserApiKeyCache, - key: str, - value: object, - *, - model_type: type[LiteLLM_TeamMembership] | None, - ttl: float | None, - write_epoch: int, -) -> None: - redis_cache: Final = user_api_key_cache.redis_cache - if redis_cache is None or _membership_write_epoch(key) != write_epoch: - return - payload: Final[object] = CacheCodec.serialize(value, model_type=model_type) - try: - if ttl is None: - await redis_cache.async_set_cache(key, payload) - else: - await redis_cache.async_set_cache(key, payload, ttl=ttl) - if _membership_write_epoch(key) != write_epoch: - await redis_cache.async_delete_cache(key) - except Exception: - return - - -async def _populate_team_membership_cache( - user_api_key_cache: UserApiKeyCache, - key: str, - value: object, - *, - model_type: type[LiteLLM_TeamMembership] | None = None, - ttl: float | None = None, -) -> None: - write_epoch: Final = _membership_write_epoch(key) - await _set_team_membership_l1( - user_api_key_cache, - key, - value, - model_type=model_type, - ttl=ttl, - ) - if _membership_write_epoch(key) != write_epoch: - user_api_key_cache.in_memory_cache_for(key).delete_cache(key) - return - - asyncio.create_task( - _replicate_team_membership_to_redis( - user_api_key_cache, - key, - value, - model_type=model_type, - ttl=ttl, - write_epoch=write_epoch, - ) - ) + return cached_membership if cached_membership is not None else _TEAM_MEMBERSHIP_CACHE_MISS @log_db_metrics @@ -2283,7 +2197,7 @@ async def _fetch_team_membership_from_db( parent_otel_span: Span | None = None, proxy_logging_obj: ProxyLogging | None = None, ) -> LiteLLM_TeamMembership | None: - """Prisma read + L1 populate. Decorated so cache hits on ``get_team_membership`` are not postgres spans.""" + """Prisma read + cache populate. Decorated so cache hits on ``get_team_membership`` are not postgres spans.""" _ = parent_otel_span, proxy_logging_obj response: Final = await _dictable_table(TeamMembershipRepository(prisma_client)).find_unique( where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, @@ -2291,19 +2205,17 @@ async def _fetch_team_membership_from_db( ) _key: Final = team_membership_reservation_cache_key(user_id=user_id, team_id=team_id) if response is None: - await _populate_team_membership_cache( - user_api_key_cache, - _key, - NO_TEAM_MEMBERSHIP_SENTINEL, + await user_api_key_cache.async_set_cache( + key=_key, + value=NO_TEAM_MEMBERSHIP_SENTINEL, ttl=get_management_object_ttl(user_api_key_cache), ) return None membership: Final = LiteLLM_TeamMembership.model_validate(response.dict()) - await _populate_team_membership_cache( - user_api_key_cache, - _key, - membership, + await user_api_key_cache.async_set_cache( + key=_key, + value=membership, model_type=LiteLLM_TeamMembership, ) return membership @@ -2321,7 +2233,7 @@ async def _load_team_membership_on_cache_miss( try: redis_cached: Final[object] = await user_api_key_cache.async_get_cache(key=cache_key) redis_membership: Final = _membership_from_cached_payload(redis_cached) - if not isinstance(redis_membership, TeamMembershipCacheMiss): + if not isinstance(redis_membership, _TeamMembershipCacheMiss): return redis_membership return await _fetch_team_membership_from_db( @@ -2332,18 +2244,9 @@ async def _load_team_membership_on_cache_miss( parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, ) - except HTTPException: - raise - except Exception as e: - verbose_proxy_logger.exception( - "Error getting team membership for user_id: %s, team_id: %s", - user_id, - team_id, - ) - raise HTTPException( - status_code=status.HTTP_503_SERVICE_UNAVAILABLE, - detail="Failed to load team membership", - ) from e + except Exception: + verbose_proxy_logger.exception("Error getting team membership") + return None async def get_team_membership( @@ -2366,7 +2269,7 @@ async def get_team_membership( l1_cached: Final[object] = await user_api_key_cache.async_get_cache(key=_key, local_only=True) l1_membership: Final = _membership_from_cached_payload(l1_cached) - if not isinstance(l1_membership, TeamMembershipCacheMiss): + if not isinstance(l1_membership, _TeamMembershipCacheMiss): return l1_membership inflight: Final[object] = _team_membership_inflight.get(_key) @@ -2375,9 +2278,6 @@ async def get_team_membership( if prisma_client is None: raise Exception("No db connected") - prisma: Final[object] = prisma_client - if isinstance(prisma, str): - return None task: Final = asyncio.ensure_future( _load_team_membership_on_cache_miss( @@ -2908,7 +2808,6 @@ async def invalidate_team_member_spend_state( ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS, ) - _bump_membership_write_epoch(team_membership_reservation_cache_key(user_id=user_id, team_id=team_id)) await evict_and_broadcast( cache_keys=( team_membership_auth_cache_key(team_id=team_id, user_id=user_id), diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 29f6c3861f6..04e41608b26 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -6494,52 +6494,6 @@ async def test_get_team_membership_coalesces_parallel_db_fetches(): mock_prisma_client.db.litellm_teammembership.find_unique.assert_awaited_once() -@pytest.mark.asyncio -async def test_get_team_membership_returns_before_redis_set_completes(): - """Auth must not wait on DualCache Redis SET; L1 is enough for the next lookup.""" - from litellm.proxy.auth.auth_checks import get_team_membership - - membership_row = MagicMock() - membership_row.dict = lambda: {"user_id": "u-redis", "team_id": "t-redis", "spend": 2.0} - - hang_redis_set = asyncio.Event() - - async def _hanging_redis_set(*args, **kwargs): - await hang_redis_set.wait() - - redis_cache = MagicMock() - redis_cache.async_get_cache = AsyncMock(return_value=None) - redis_cache.async_set_cache = AsyncMock(side_effect=_hanging_redis_set) - - cache = UserApiKeyCache(redis_cache=redis_cache) - mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row) - - first = await asyncio.wait_for( - get_team_membership( - user_id="u-redis", - team_id="t-redis", - prisma_client=mock_prisma_client, - user_api_key_cache=cache, - ), - timeout=0.5, - ) - second = await get_team_membership( - user_id="u-redis", - team_id="t-redis", - prisma_client=mock_prisma_client, - user_api_key_cache=cache, - ) - - hang_redis_set.set() - await asyncio.sleep(0) - - assert first is not None and second is not None - assert first.spend == 2.0 - assert second.spend == 2.0 - mock_prisma_client.db.litellm_teammembership.find_unique.assert_awaited_once() - - @pytest.mark.asyncio async def test_common_checks_calls_get_team_membership_once_per_request(): """Model-access, attribution, and member-budget must reuse one membership load.""" @@ -6588,33 +6542,85 @@ async def test_common_checks_calls_get_team_membership_once_per_request(): @pytest.mark.asyncio -async def test_get_team_membership_db_error_raises_503_not_none(): - """A Prisma failure must fail closed as 503, not look like a missing membership row.""" - from fastapi import HTTPException +async def test_common_checks_skips_membership_load_when_no_check_reads_it(): + """A management route has no model and no budget gate, so the membership row is never loaded.""" + from fastapi import Request + from litellm.proxy.auth.auth_checks import common_checks + + team = LiteLLM_TeamTable(team_id="t-lazy") + token = UserAPIKeyAuth(token="k-lazy", user_id="u-lazy", team_id="t-lazy") + + with ( + patch( # test-quality-ok: common_checks imports prisma_client from proxy_server + "litellm.proxy.proxy_server.prisma_client", MagicMock() + ), + patch( # test-quality-ok: common_checks imports user_api_key_cache from proxy_server + "litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache() + ), + patch( # test-quality-ok: counts membership loads; common_checks has no membership seam + "litellm.proxy.auth.auth_checks.get_team_membership", + new_callable=AsyncMock, + ) as load_membership, + ): + result = await common_checks( + request_body={}, + team_object=team, + user_object=LiteLLM_UserTable(user_id="u-lazy"), + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/key/info", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=token, + request=MagicMock(spec=Request), + ) + + assert result is True + load_membership.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_get_team_membership_db_error_returns_none_and_retries_next_call(): + """A Prisma failure reads as no membership, caches nothing, and the next call hits the DB again.""" from litellm.proxy.auth.auth_checks import get_team_membership from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key + membership_row = MagicMock() + membership_row.dict = lambda: {"user_id": "u-fail", "team_id": "t-fail", "spend": 1.0} mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(side_effect=RuntimeError("db down")) + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock( + side_effect=[RuntimeError("db down"), membership_row] + ) cache = UserApiKeyCache() - with pytest.raises(HTTPException) as exc: - await get_team_membership( - user_id="u-fail", - team_id="t-fail", - prisma_client=mock_prisma_client, - user_api_key_cache=cache, - ) + failed = await get_team_membership( + user_id="u-fail", + team_id="t-fail", + prisma_client=mock_prisma_client, + user_api_key_cache=cache, + ) + cached_after_failure = await cache.async_get_cache( + key=team_membership_reservation_cache_key(user_id="u-fail", team_id="t-fail") + ) + recovered = await get_team_membership( + user_id="u-fail", + team_id="t-fail", + prisma_client=mock_prisma_client, + user_api_key_cache=cache, + ) - cached = await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-fail", team_id="t-fail")) - assert exc.value.status_code == status.HTTP_503_SERVICE_UNAVAILABLE - assert cached is None + assert failed is None + assert cached_after_failure is None + assert recovered is not None + assert recovered.user_id == "u-fail" + assert mock_prisma_client.db.litellm_teammembership.find_unique.await_count == 2 @pytest.mark.asyncio async def test_get_team_membership_string_prisma_client_returns_none(): - """Unit tests stub prisma_client as a string; that is not a lookup failure and must not 503.""" + """Unit tests stub prisma_client as a string; the lookup fails and reads as no membership.""" from litellm.proxy.auth.auth_checks import get_team_membership result = await get_team_membership( @@ -6626,59 +6632,6 @@ async def test_get_team_membership_string_prisma_client_returns_none(): assert result is None -@pytest.mark.asyncio -async def test_common_checks_does_not_skip_member_limits_when_membership_lookup_fails(): - """common_checks must not mark membership loaded-absent after a lookup error.""" - from fastapi import HTTPException, Request - - from litellm.proxy.auth.auth_checks import common_checks - - team = LiteLLM_TeamTable(team_id="t-fail-closed") - token = UserAPIKeyAuth( - token="k-fail-closed", - user_id="u-fail-closed", - team_id="t-fail-closed", - models=["gpt-4o-mini"], - ) - lookup_error = HTTPException( - status_code=status.HTTP_503_SERVICE_UNAVAILABLE, - detail="Failed to load team membership", - ) - - with ( - patch( # test-quality-ok: common_checks imports prisma_client from proxy_server - "litellm.proxy.proxy_server.prisma_client", MagicMock() - ), - patch( # test-quality-ok: common_checks imports user_api_key_cache from proxy_server - "litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache() - ), - patch( # test-quality-ok: injects membership lookup failure; common_checks has no seam - "litellm.proxy.auth.auth_checks.get_team_membership", - new_callable=AsyncMock, - side_effect=lookup_error, - ), - patch( # test-quality-ok: common_checks imports get_current_spend locally - "litellm.proxy.proxy_server.get_current_spend", new_callable=AsyncMock, return_value=0.0 - ), - ): - with pytest.raises(HTTPException) as exc: - await common_checks( - request_body={"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hi"}]}, - team_object=team, - user_object=LiteLLM_UserTable(user_id="u-fail-closed"), - end_user_object=None, - global_proxy_spend=None, - general_settings={}, - route="/chat/completions", - llm_router=None, - proxy_logging_obj=MagicMock(), - valid_token=token, - request=MagicMock(spec=Request), - ) - - assert exc.value.status_code == status.HTTP_503_SERVICE_UNAVAILABLE - - @pytest.mark.asyncio async def test_get_team_membership_waiter_cancel_does_not_cancel_shared_load(): """Cancelling one coalesced waiter must not cancel the shared Prisma load.""" @@ -6721,51 +6674,6 @@ async def test_get_team_membership_waiter_cancel_does_not_cancel_shared_load(): mock_prisma_client.db.litellm_teammembership.find_unique.assert_awaited_once() -@pytest.mark.asyncio -async def test_stale_membership_redis_replicate_does_not_restore_after_invalidate(): - """A delayed Redis SET must not rewrite L1 or leave Redis holding membership after invalidation.""" - from litellm.proxy.auth.auth_checks import get_team_membership, invalidate_team_member_spend_state - from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key - - membership_row = MagicMock() - membership_row.dict = lambda: {"user_id": "u-stale", "team_id": "t-stale", "spend": 9.0} - - hang_redis_set = asyncio.Event() - - async def _hanging_redis_set(*args, **kwargs): - await hang_redis_set.wait() - - redis_cache = MagicMock() - redis_cache.async_get_cache = AsyncMock(return_value=None) - redis_cache.async_set_cache = AsyncMock(side_effect=_hanging_redis_set) - redis_cache.async_delete_cache = AsyncMock() - redis_cache.delete_cache = MagicMock() - - cache = UserApiKeyCache(redis_cache=redis_cache) - mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row) - - loaded = await get_team_membership( - user_id="u-stale", - team_id="t-stale", - prisma_client=mock_prisma_client, - user_api_key_cache=cache, - ) - cache_key = team_membership_reservation_cache_key(user_id="u-stale", team_id="t-stale") - await invalidate_team_member_spend_state(user_id="u-stale", team_id="t-stale", user_api_key_cache=cache) - after_invalidate = await cache.async_get_cache(key=cache_key, local_only=True) - - hang_redis_set.set() - await asyncio.sleep(0) - await asyncio.sleep(0) - after_replicate = await cache.async_get_cache(key=cache_key, local_only=True) - - assert loaded is not None - assert after_invalidate is None - assert after_replicate is None - redis_cache.async_delete_cache.assert_awaited() - - @pytest.mark.asyncio async def test_invalidate_team_member_spend_state_evicts_the_negative_cache_sentinel(): """ From cb5d901774379b8aef4b946c84cb722877d5ba91 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 09:53:39 +0000 Subject: [PATCH 076/112] fix(bedrock/realtime): propagate deferred Nova Sonic stream failures to the router Bedrock realtime caught every exception inside both forwarding tasks and gathered them with return_exceptions=True, so a provider failure surfacing after the websocket handshake (lazy duplex stream: 503/429/validation only show up on await_output or the input publisher) made async_realtime return normally and the router recorded a success instead of running fallbacks and cooldown accounting. session.updated is now acked only after Bedrock is ready, provider failures escape as BedrockError with the AWS status code, a failure after the client disconnected is not reported as a provider failure, and a fallback attempt on the same websocket replays the pending session.update instead of emitting a second session.created. Resolves LIT-6484 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/bedrock/realtime/handler.py | 451 +++++++++++------- .../realtime/test_bedrock_realtime_handler.py | 166 ++++++- 2 files changed, 453 insertions(+), 164 deletions(-) diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index ca2370303f2..ca5b87d1700 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -7,7 +7,9 @@ This uses aws_sdk_bedrock_runtime for bidirectional streaming with Nova Sonic. import asyncio import contextlib import json -from collections.abc import AsyncIterator, Mapping +from collections.abc import AsyncIterator, Mapping, MutableMapping +from dataclasses import dataclass +from types import MappingProxyType from typing import Final, Protocol from pydantic import JsonValue, TypeAdapter @@ -28,6 +30,43 @@ from .transformation import BedrockRealtimeConfig _CLIENT_MODALITIES_ADAPTER: Final[TypeAdapter["list[str] | None"]] = TypeAdapter(list[str] | None) _CLIENT_MESSAGE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) +_EMPTY_JSON_OBJECT: Final[Mapping[str, JsonValue]] = MappingProxyType({}) +_PENDING_UPDATE_KEY: Final = "litellm.bedrock_realtime.pending_session_update" +_COMMITTED_KEY: Final = "litellm.bedrock_realtime.session_committed" + +_BEDROCK_STREAM_ERROR_STATUS: Final[Mapping[str, int]] = MappingProxyType( + { + "AccessDeniedException": 403, + "ConflictException": 400, + "InternalServerException": 500, + "ModelErrorException": 424, + "ModelNotReadyException": 429, + "ModelStreamErrorException": 424, + "ModelTimeoutException": 408, + "ResourceNotFoundException": 404, + "ServiceQuotaExceededException": 400, + "ServiceUnavailableException": 503, + "ThrottlingException": 429, + "ValidationException": 400, + } +) + + +def _as_bedrock_error(error: BaseException) -> BaseException: + status_code: Final = _BEDROCK_STREAM_ERROR_STATUS.get(type(error).__name__) + if status_code is None: + return error + return BedrockError(status_code=status_code, message=f"{type(error).__name__}: {error}") + + +class _BedrockForwardingFailed(Exception): + """The Bedrock output stream failed after ``logged_events`` were already forwarded to the client.""" + + def __init__(self, cause: BaseException, logged_events: tuple[OpenAIRealtimeEvents, ...]) -> None: + super().__init__(str(cause)) + self.cause: Final = cause + self.logged_events: Final = logged_events + def _json_dict(value: JsonValue) -> dict[str, JsonValue]: return value if isinstance(value, dict) else {} @@ -51,6 +90,9 @@ def _should_log_event(openai_message: Mapping[str, object]) -> bool: class RealtimeClientWebSocket(Protocol): """The client-facing websocket surface the realtime bridge talks to.""" + @property + def scope(self) -> MutableMapping[str, object]: ... # mutable-ok: the ASGI scope is the per-connection state store + async def receive_text(self) -> str: ... async def send_text(self, data: str) -> None: ... @@ -85,6 +127,71 @@ class BedrockBidirectionalStream(Protocol): async def await_output(self) -> tuple[object, BedrockOutputStream]: ... +@dataclass(frozen=True, slots=True) +class _BridgeOutcome: + logged_events: tuple[OpenAIRealtimeEvents, ...] + provider_failure: BaseException | None + client_disconnected: bool + + +async def _client_messages(client_ws: RealtimeClientWebSocket, initial_message: str | None) -> AsyncIterator[str]: + if initial_message is not None: + yield initial_message + while True: + try: + yield await client_ws.receive_text() + except Exception as e: # noqa: BLE001 # any receive failure means the client is gone + verbose_proxy_logger.debug("Client to Bedrock forwarding ended: %s", e, exc_info=True) + return + + +def _take_pending_session_update( + scope: MutableMapping[str, object], # mutable-ok: the ASGI scope is the per-connection state store +) -> str | None: + """A fallback attempt on the same websocket replays the session.update the failed attempt never acked.""" + if scope.get(_COMMITTED_KEY) is True: + raise BedrockError( + status_code=409, + message="Bedrock realtime session already committed to a provider stream; it cannot be replayed", + ) + pending: Final = scope.pop(_PENDING_UPDATE_KEY, None) # rebind-ok: the ASGI scope outlives this attempt + return pending if isinstance(pending, str) else None + + +def _parse_client_message(message: str) -> Mapping[str, JsonValue]: + try: + return _json_dict(_CLIENT_MESSAGE_ADAPTER.validate_json(message)) + except ValueError: + return _EMPTY_JSON_OBJECT + + +async def _ack_session_update( + client_ws: RealtimeClientWebSocket, + bedrock_stream: BedrockBidirectionalStream, + transformation_config: BedrockRealtimeConfig, + model: str, + logging_obj: LiteLLMLogging | None, + parsed_client_message: Mapping[str, JsonValue], +) -> bool: + """Ack the client's session.update once Bedrock accepted the stream; False means the client is gone.""" + await bedrock_stream.await_output() + client_ws.scope.pop(_PENDING_UPDATE_KEY, None) + client_ws.scope[_COMMITTED_KEY] = True # rebind-ok: the ASGI scope outlives this attempt + if logging_obj is None: + return True + requested_modalities: Final = _CLIENT_MODALITIES_ADAPTER.validate_python( + _json_dict(parsed_client_message.get("session")).get("modalities") + ) + try: + await client_ws.send_text( + json.dumps(transformation_config.session_updated_event(model, logging_obj, requested_modalities)) + ) + except Exception as e: # noqa: BLE001 # any send failure means the client is gone + verbose_proxy_logger.debug("Client to Bedrock forwarding ended: %s", e, exc_info=True) + return False + return True + + class BedrockRealtime(BaseAWSLLM): """Handler for Bedrock Nova Sonic realtime speech-to-speech API.""" @@ -132,6 +239,8 @@ class BedrockRealtime(BaseAWSLLM): except ImportError: raise ImportError("Missing aws_sdk_bedrock_runtime. Install with: pip install aws-sdk-bedrock-runtime") + pending_session_update: Final = _take_pending_session_update(websocket.scope) + # Get AWS region if aws_region_name is None: optional_params: Final = { @@ -190,90 +299,118 @@ class BedrockRealtime(BaseAWSLLM): transformation_config: Final = BedrockRealtimeConfig() - try: - # Initialize the bidirectional stream - bedrock_stream: Final = await open_bidirectional_stream() + bedrock_stream: Final = await open_bidirectional_stream() - verbose_proxy_logger.debug("Bedrock Realtime: Bidirectional stream established") + verbose_proxy_logger.debug("Bedrock Realtime: Bidirectional stream established") + if pending_session_update is None: await websocket.send_text(json.dumps(transformation_config.session_created_event(model, logging_obj))) verbose_proxy_logger.debug("Bedrock Realtime: sent session.created to client on connect") - # Track state for transformation - session_state: Final[RealtimeResponseTransformInput] = { - "current_output_item_id": None, - "current_response_id": None, - "current_conversation_id": None, - "current_delta_chunks": None, - "current_item_chunks": None, - "current_delta_type": None, - "session_configuration_request": None, - } + # Track state for transformation + session_state: Final[RealtimeResponseTransformInput] = { + "current_output_item_id": None, + "current_response_id": None, + "current_conversation_id": None, + "current_delta_chunks": None, + "current_item_chunks": None, + "current_delta_type": None, + "session_configuration_request": None, + } - # Create tasks for bidirectional forwarding - client_to_bedrock_task: Final = asyncio.create_task( - self._forward_client_to_bedrock( - websocket, - bedrock_stream, - transformation_config, - model, - session_state, - logging_obj, + outcome: Final = await self._bridge( + websocket, + bedrock_stream, + transformation_config, + model, + session_state, + logging_obj, + initial_message=pending_session_update, + ) + + logged_events: Final = ( + *outcome.logged_events, + *( + leftover_event + for leftover_event in transformation_config.leftover_usage_done_events() + if _should_log_event(leftover_event) + ), + ) + if logged_events: + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( + logging_obj.dispatch_success_handlers( + list(logged_events), # mutable-ok: realtime spend logging requires a list result + prefer_async_handlers=True, ) ) - async def forward_bedrock_and_collect_logged_events() -> tuple[OpenAIRealtimeEvents, ...]: - return tuple( - [ - event - async for event in self._forward_bedrock_to_client( - bedrock_stream, - websocket, - transformation_config, - model, - logging_obj, - session_state, - ) - ] - ) - - bedrock_to_client_task: Final = asyncio.create_task(forward_bedrock_and_collect_logged_events()) - - # Wait for both tasks to complete - await asyncio.gather( - client_to_bedrock_task, - bedrock_to_client_task, - return_exceptions=True, + if outcome.provider_failure is None: + return + if outcome.client_disconnected: + verbose_proxy_logger.debug( + "Bedrock Realtime: stream failed after the client disconnected: %s", outcome.provider_failure ) + return + verbose_proxy_logger.error( + "Bedrock Realtime: provider stream failed: %s", _redact_string(str(outcome.provider_failure)) + ) + raise _as_bedrock_error(outcome.provider_failure) from outcome.provider_failure - forwarded_logged_events: Final = ( - bedrock_to_client_task.result() - if not bedrock_to_client_task.cancelled() and bedrock_to_client_task.exception() is None - else () - ) - logged_events: Final = ( - *forwarded_logged_events, - *( - leftover_event - for leftover_event in transformation_config.leftover_usage_done_events() - if _should_log_event(leftover_event) - ), - ) - if logged_events: - GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( - logging_obj.dispatch_success_handlers( - list(logged_events), # mutable-ok: realtime spend logging requires a list result - prefer_async_handlers=True, - ) - ) + async def _bridge( + self, + websocket: RealtimeClientWebSocket, + bedrock_stream: BedrockBidirectionalStream, + transformation_config: BedrockRealtimeConfig, + model: str, + session_state: RealtimeResponseTransformInput, + logging_obj: LiteLLMLogging, + initial_message: str | None, + ) -> _BridgeOutcome: + """Run both forwarding directions until the client leaves or either side fails.""" - except Exception as e: - verbose_proxy_logger.exception("Error in BedrockRealtime.async_realtime: %s", e) + async def collect_logged_events() -> tuple[OpenAIRealtimeEvents, ...]: + logged: Final[list[OpenAIRealtimeEvents]] = [] # mutable-ok: partial events are still logged try: - await websocket.close(code=1011, reason=_redact_string(f"Internal error: {e}")) - except Exception: - pass - raise + async for event in self._forward_bedrock_to_client( + bedrock_stream, websocket, transformation_config, model, logging_obj, session_state + ): + logged.append(event) + except Exception as e: + raise _BedrockForwardingFailed(e, tuple(logged)) from e + return tuple(logged) + + client_task: Final = asyncio.create_task( + self._forward_client_to_bedrock( + websocket, bedrock_stream, transformation_config, model, session_state, logging_obj, initial_message + ) + ) + bedrock_task: Final = asyncio.create_task(collect_logged_events()) + + await asyncio.wait((client_task, bedrock_task), return_when=asyncio.FIRST_EXCEPTION) + client_disconnected: Final = ( + client_task.done() and not client_task.cancelled() and client_task.exception() is None + ) + client_task.cancel() + bedrock_task.cancel() + client_outcome, bedrock_outcome = await asyncio.gather(client_task, bedrock_task, return_exceptions=True) + + return _BridgeOutcome( + logged_events=( + bedrock_outcome.logged_events + if isinstance(bedrock_outcome, _BedrockForwardingFailed) + else bedrock_outcome + if isinstance(bedrock_outcome, tuple) + else () + ), + provider_failure=( + client_outcome + if isinstance(client_outcome, Exception) + else bedrock_outcome.cause + if isinstance(bedrock_outcome, _BedrockForwardingFailed) + else None + ), + client_disconnected=client_disconnected, + ) async def _forward_client_to_bedrock( self, @@ -283,8 +420,12 @@ class BedrockRealtime(BaseAWSLLM): model: str, session_state: RealtimeResponseTransformInput, logging_obj: LiteLLMLogging | None = None, - ): - """Forward messages from client WebSocket to Bedrock stream.""" + initial_message: str | None = None, + ) -> None: + """Forward messages from client WebSocket to Bedrock stream. + + Returns once the client is gone; provider failures (input stream or readiness) propagate to the caller. + """ from aws_sdk_bedrock_runtime.models import ( BidirectionalInputPayloadPart, InvokeModelWithBidirectionalStreamInputChunk, @@ -299,41 +440,26 @@ class BedrockRealtime(BaseAWSLLM): verbose_proxy_logger.debug("Bedrock Realtime: Sent to Bedrock: %s", bedrock_message[:200]) try: - while True: - # Receive message from client - message = await client_ws.receive_text() + async for message in _client_messages(client_ws, initial_message): verbose_proxy_logger.debug("Bedrock Realtime: Received from client: %s", message[:200]) + parsed_client_message = _parse_client_message(message) + is_session_update = _json_str(parsed_client_message.get("type")) == "session.update" + if is_session_update: + client_ws.scope[_PENDING_UPDATE_KEY] = message # rebind-ok: scope outlives the attempt - # Transform OpenAI format to Bedrock format transformed_messages = transformation_config.transform_realtime_request( message=message, model=model, session_configuration_request=session_state.get("session_configuration_request"), ) - - # Send transformed messages to Bedrock for bedrock_message in transformed_messages: await send_to_bedrock(bedrock_message) - if logging_obj is not None: - client_message_type: str | None = None - requested_modalities: list[str] | None = None - with contextlib.suppress(Exception): - parsed_client_message = _json_dict(_CLIENT_MESSAGE_ADAPTER.validate_json(message)) - client_message_type = _json_str(parsed_client_message.get("type")) - if client_message_type == "session.update": - requested_modalities = _CLIENT_MODALITIES_ADAPTER.validate_python( - _json_dict(parsed_client_message.get("session")).get("modalities") - ) - if client_message_type == "session.update": - await client_ws.send_text( - json.dumps( - transformation_config.session_updated_event(model, logging_obj, requested_modalities) - ) - ) - - except Exception as e: - verbose_proxy_logger.debug("Client to Bedrock forwarding ended: %s", e, exc_info=True) + if is_session_update and not await _ack_session_update( + client_ws, bedrock_stream, transformation_config, model, logging_obj, parsed_client_message + ): + break + finally: for close_message in transformation_config.session_close_messages(): with contextlib.suppress(Exception): await send_to_bedrock(close_message) @@ -349,68 +475,71 @@ class BedrockRealtime(BaseAWSLLM): logging_obj: LiteLLMLogging, session_state: RealtimeResponseTransformInput, ) -> AsyncIterator[OpenAIRealtimeEvents]: - """Forward messages from Bedrock to the client, yielding the ones to record for spend logging.""" - try: - while True: - # Receive from Bedrock - output = await bedrock_stream.await_output() - result = await output[1].receive() + """Forward messages from Bedrock to the client, yielding the ones to record for spend logging. - if result is None: - verbose_proxy_logger.debug("Bedrock Realtime: Bedrock stream ended") - break + Provider failures propagate to the caller; the client websocket is only closed on a normal stream end. + """ - payload_bytes = result.value.bytes_ if result.value else None - if payload_bytes: - bedrock_response = payload_bytes.decode("utf-8") - verbose_proxy_logger.debug("Bedrock Realtime: Received from Bedrock: %s", bedrock_response[:200]) - - # Transform Bedrock format to OpenAI format - realtime_response_transform_input: RealtimeResponseTransformInput = { - "current_output_item_id": session_state.get("current_output_item_id"), - "current_response_id": session_state.get("current_response_id"), - "current_conversation_id": session_state.get("current_conversation_id"), - "current_delta_chunks": session_state.get("current_delta_chunks"), - "current_item_chunks": session_state.get("current_item_chunks"), - "current_delta_type": session_state.get("current_delta_type"), - "session_configuration_request": session_state.get("session_configuration_request"), - } - - transformed_response = transformation_config.transform_realtime_response( - message=bedrock_response, - model=model, - logging_obj=logging_obj, - realtime_response_transform_input=realtime_response_transform_input, - ) - - # Update session state - session_state.update( - { - "current_output_item_id": transformed_response.get("current_output_item_id"), - "current_response_id": transformed_response.get("current_response_id"), - "current_conversation_id": transformed_response.get("current_conversation_id"), - "current_delta_chunks": transformed_response.get("current_delta_chunks"), - "current_item_chunks": transformed_response.get("current_item_chunks"), - "current_delta_type": transformed_response.get("current_delta_type"), - "session_configuration_request": transformed_response.get("session_configuration_request"), - } - ) - - # Send transformed messages to client - response_value = transformed_response["response"] - openai_messages = response_value if isinstance(response_value, list) else (response_value,) - for openai_message in openai_messages: - message_json = json.dumps(openai_message) - await client_ws.send_text(message_json) - verbose_proxy_logger.debug("Bedrock Realtime: Sent to client: %s", message_json[:200]) - if _should_log_event(openai_message): - yield openai_message - - except Exception as e: - verbose_proxy_logger.debug("Bedrock to client forwarding ended: %s", e, exc_info=True) - finally: - # Close the client WebSocket + async def send_to_client(message_json: str) -> bool: try: - await client_ws.close() - except Exception: - pass + await client_ws.send_text(message_json) + except Exception as e: # noqa: BLE001 # any send failure means the client is gone + verbose_proxy_logger.debug("Bedrock to client forwarding ended: %s", e, exc_info=True) + return False + verbose_proxy_logger.debug("Bedrock Realtime: Sent to client: %s", message_json[:200]) + return True + + output: Final = await bedrock_stream.await_output() + while True: + result = await output[1].receive() + + if result is None: + verbose_proxy_logger.debug("Bedrock Realtime: Bedrock stream ended") + with contextlib.suppress(Exception): + await client_ws.close() + return + + payload_bytes = result.value.bytes_ if result.value else None + if payload_bytes: + bedrock_response = payload_bytes.decode("utf-8") + verbose_proxy_logger.debug("Bedrock Realtime: Received from Bedrock: %s", bedrock_response[:200]) + + # Transform Bedrock format to OpenAI format + realtime_response_transform_input: RealtimeResponseTransformInput = { + "current_output_item_id": session_state.get("current_output_item_id"), + "current_response_id": session_state.get("current_response_id"), + "current_conversation_id": session_state.get("current_conversation_id"), + "current_delta_chunks": session_state.get("current_delta_chunks"), + "current_item_chunks": session_state.get("current_item_chunks"), + "current_delta_type": session_state.get("current_delta_type"), + "session_configuration_request": session_state.get("session_configuration_request"), + } + + transformed_response = transformation_config.transform_realtime_response( + message=bedrock_response, + model=model, + logging_obj=logging_obj, + realtime_response_transform_input=realtime_response_transform_input, + ) + + # Update session state + session_state.update( + { + "current_output_item_id": transformed_response.get("current_output_item_id"), + "current_response_id": transformed_response.get("current_response_id"), + "current_conversation_id": transformed_response.get("current_conversation_id"), + "current_delta_chunks": transformed_response.get("current_delta_chunks"), + "current_item_chunks": transformed_response.get("current_item_chunks"), + "current_delta_type": transformed_response.get("current_delta_type"), + "session_configuration_request": transformed_response.get("session_configuration_request"), + } + ) + + # Send transformed messages to client + response_value = transformed_response["response"] + openai_messages = response_value if isinstance(response_value, list) else (response_value,) + for openai_message in openai_messages: + if not await send_to_client(json.dumps(openai_message)): + return + if _should_log_event(openai_message): + yield openai_message diff --git a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py index 0ea5b7ad4a1..0439a5090f5 100644 --- a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py +++ b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py @@ -1,3 +1,4 @@ +import asyncio import json import sys import types @@ -51,6 +52,27 @@ class FakeBedrockStream: def __init__(self, input_stream=None): self.input_stream = input_stream if input_stream is not None else FakeInputStream() + async def await_output(self): + return (None, EndedBedrockReceiver()) + + +class ServiceUnavailableException(Exception): + """Named like the modeled AWS SDK error so the handler maps it to HTTP 503""" + + +class ModelStreamErrorException(Exception): + """Named like the modeled AWS SDK error so the handler maps it to HTTP 424""" + + +class UnavailableBedrockStream: + """Lazy duplex stream whose HTTP response only fails once the output is awaited""" + + def __init__(self): + self.input_stream = FakeInputStream() + + async def await_output(self): + raise ServiceUnavailableException("fault injected: Bedrock realtime unavailable") + class FakeLogging: def __init__(self, trace_id="trace-nova-sonic"): @@ -61,6 +83,7 @@ class DisconnectingClientWS: def __init__(self, messages): self._messages = list(messages) self.sent_to_client = [] + self.scope = {} async def receive_text(self): if self._messages: @@ -93,6 +116,7 @@ class RealtimeClientWS: def __init__(self): self.closed = False self.sent_to_client = [] + self.scope = {} async def receive_text(self): raise RuntimeError("client disconnected") @@ -104,6 +128,25 @@ class RealtimeClientWS: self.closed = True +class ConnectedClientWS(RealtimeClientWS): + """Client that sends its scripted messages and then stays connected until the server closes it""" + + def __init__(self, messages): + super().__init__() + self._messages = list(messages) + self._closed_event = asyncio.Event() + + async def receive_text(self): + if self._messages: + return self._messages.pop(0) + await self._closed_event.wait() + raise RuntimeError("client disconnected") + + async def close(self, code=None, reason=None): + self.closed = True + self._closed_event.set() + + class ScriptedBedrockReceiver: def __init__(self, payloads): self._payloads = list(payloads) @@ -115,10 +158,20 @@ class ScriptedBedrockReceiver: return SimpleNamespace(value=SimpleNamespace(bytes_=payload.encode("utf-8"))) +class BreakingBedrockReceiver(ScriptedBedrockReceiver): + """Delivers its payloads, then the provider stream breaks instead of ending normally""" + + async def receive(self): + if not self._payloads: + await asyncio.sleep(0) + raise ModelStreamErrorException("Nova Sonic stream broke") + return await super().receive() + + class ScriptedBedrockStream: - def __init__(self, payloads): + def __init__(self, payloads, receiver_type=ScriptedBedrockReceiver): self.input_stream = FakeInputStream() - self._receiver = ScriptedBedrockReceiver(payloads) + self._receiver = receiver_type(payloads) async def await_output(self): return (None, self._receiver) @@ -163,6 +216,8 @@ def stub_aws_sdk_client(monkeypatch): async def invoke_model_with_bidirectional_stream(self, operation_input): captured["operation_input"] = operation_input + if captured.get("streams"): + return captured["streams"].pop(0) return ScriptedBedrockStream(captured.get("scripted_payloads", [])) package = types.ModuleType("aws_sdk_bedrock_runtime") @@ -263,7 +318,8 @@ class TestBedrockRealtimeHandler: [json.dumps({"type": "session.update", "session": {"instructions": "You are helpful."}})] ) - await handler._forward_client_to_bedrock(client_ws, stream, config, "amazon.nova-sonic-v1:0", {}) + with pytest.raises(RuntimeError, match="bedrock send failed"): + await handler._forward_client_to_bedrock(client_ws, stream, config, "amazon.nova-sonic-v1:0", {}) assert stream.input_stream.closed @@ -464,6 +520,110 @@ class TestBedrockRealtimeSessionLifecycle: assert client_ws.sent_to_client == [] +class TestBedrockRealtimeProviderFailurePropagation: + """Deferred Nova Sonic failures must escape async_realtime so the router can fall back / cool down (LIT-6484)""" + + SESSION_UPDATE = json.dumps({"type": "session.update", "session": {"instructions": "hi", "modalities": ["text"]}}) + AWS_PARAMS = {"aws_region_name": "us-east-1", "aws_access_key_id": "k", "aws_secret_access_key": "s"} + + @pytest.mark.asyncio + async def test_readiness_failure_escapes_and_fallback_replays_session_update(self, stub_aws_sdk_client): + handler = BedrockRealtime() + websocket = ConnectedClientWS([self.SESSION_UPDATE]) + healthy_stream = ScriptedBedrockStream([]) + stub_aws_sdk_client["streams"] = [UnavailableBedrockStream(), healthy_stream] + + with pytest.raises(BedrockError) as failure: + await handler.async_realtime( + model="amazon.nova-sonic-v1:0", websocket=websocket, logging_obj=FakeLogging(), **self.AWS_PARAMS + ) + + assert failure.value.status_code == 503 + assert [json.loads(m)["type"] for m in websocket.sent_to_client] == ["session.created"] + assert not websocket.closed, "the proxy route owns the client-facing error event and 1011 close" + + await handler.async_realtime( + model="amazon.nova-sonic-v1:0", websocket=websocket, logging_obj=FakeLogging(), **self.AWS_PARAMS + ) + + assert [json.loads(m)["type"] for m in websocket.sent_to_client] == ["session.created", "session.updated"] + replayed = [json.loads(chunk.value.bytes_.decode("utf-8")) for chunk in healthy_stream.input_stream.sent] + assert [next(iter(event["event"])) for event in replayed][:2] == ["sessionStart", "promptStart"] + assert websocket.closed + + @pytest.mark.asyncio + async def test_mid_stream_failure_escapes_keeps_partial_spend_and_blocks_replay( + self, stub_aws_sdk_client, monkeypatch + ): + import litellm.llms.bedrock.realtime.handler as handler_module + + dispatched = {} + + class RecordingLogging(FakeLogging): + async def dispatch_success_handlers(self, result=None, prefer_async_handlers=False, **kwargs): + dispatched["events"] = result + + class RecordingLoggingWorker: + def ensure_initialized_and_enqueue(self, coro): + dispatched["coro"] = coro + + monkeypatch.setattr(handler_module, "GLOBAL_LOGGING_WORKER", RecordingLoggingWorker()) + handler = BedrockRealtime() + websocket = ConnectedClientWS([self.SESSION_UPDATE]) + stub_aws_sdk_client["streams"] = [ + ScriptedBedrockStream( + [ + json.dumps({"event": {"contentStart": {"role": "ASSISTANT", "type": "TEXT"}}}), + json.dumps({"event": {"textOutput": {"content": "Hi"}}}), + json.dumps({"event": {"contentEnd": {"stopReason": "END_TURN"}}}), + ], + receiver_type=BreakingBedrockReceiver, + ) + ] + + with pytest.raises(BedrockError) as failure: + await handler.async_realtime( + model="amazon.nova-sonic-v1:0", websocket=websocket, logging_obj=RecordingLogging(), **self.AWS_PARAMS + ) + + assert failure.value.status_code == 424 + await dispatched["coro"] + assert [event["type"] for event in dispatched["events"]] == ["response.done"] + assert "response.done" in [json.loads(m)["type"] for m in websocket.sent_to_client] + + with pytest.raises(BedrockError) as replay: + await handler.async_realtime( + model="amazon.nova-sonic-v1:0", websocket=websocket, logging_obj=RecordingLogging(), **self.AWS_PARAMS + ) + + assert replay.value.status_code == 409, "a committed session must not be silently restarted on a fallback" + + @pytest.mark.asyncio + async def test_stream_failure_after_client_disconnect_is_not_a_provider_failure(self, stub_aws_sdk_client): + stream = ScriptedBedrockStream([], receiver_type=BreakingBedrockReceiver) + stub_aws_sdk_client["streams"] = [stream] + + await BedrockRealtime().async_realtime( + model="amazon.nova-sonic-v1:0", websocket=RealtimeClientWS(), logging_obj=FakeLogging(), **self.AWS_PARAMS + ) + + assert stream.input_stream.closed + + @pytest.mark.asyncio + async def test_session_updated_is_not_sent_before_bedrock_is_ready(self, stub_aws_models): + handler = BedrockRealtime() + stream = UnavailableBedrockStream() + client_ws = DisconnectingClientWS([self.SESSION_UPDATE]) + + with pytest.raises(ServiceUnavailableException): + await handler._forward_client_to_bedrock( + client_ws, stream, BedrockRealtimeConfig(), "amazon.nova-sonic-v1:0", {}, FakeLogging() + ) + + assert client_ws.sent_to_client == [] + assert stream.input_stream.closed + + class TestBedrockRealtimeAwsAuth: """AWS auth params passed via litellm_params must reach the Smithy client config (LIT-3923 regression)""" From 5c04ec2b93eee854d6c7f36739e2cd34a226a49b Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 10:21:28 +0000 Subject: [PATCH 077/112] fix(bedrock/realtime): keep the pending session.update until a provider stream is committed Peek at the pending session.update instead of popping it, so an eager fallback failure before the bridge starts does not lose the replay for the next attempt. Move the websocket scope keys to constants.py Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 4 ++++ litellm/llms/bedrock/realtime/handler.py | 24 ++++++++++--------- .../realtime/test_bedrock_realtime_handler.py | 13 ++++++++-- 3 files changed, 28 insertions(+), 13 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 5751e6e46af..7d86080086a 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -311,6 +311,10 @@ REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS: Final = float( # RFC 6455 caps the close frame payload at 125 bytes, 2 of which carry the status code WEBSOCKET_CLOSE_REASON_MAX_BYTES: Final = 123 +# ASGI websocket scope keys the Bedrock realtime bridge uses to carry state across router fallback attempts +BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY: Final = "litellm.bedrock_realtime.pending_session_update" +BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY: Final = "litellm.bedrock_realtime.session_committed" + # SSL/TLS cipher configuration for faster handshakes # Strategy: Strongly prefer fast modern ciphers, but allow fallback to commonly supported ones # This balances performance with broad compatibility diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index ca5b87d1700..126a7ae5c3e 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -16,6 +16,10 @@ from pydantic import JsonValue, TypeAdapter import litellm from litellm._logging import _redact_string, verbose_proxy_logger +from litellm.constants import ( + BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY, + BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY, +) from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER @@ -31,8 +35,6 @@ _CLIENT_MODALITIES_ADAPTER: Final[TypeAdapter["list[str] | None"]] = TypeAdapter _CLIENT_MESSAGE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) _EMPTY_JSON_OBJECT: Final[Mapping[str, JsonValue]] = MappingProxyType({}) -_PENDING_UPDATE_KEY: Final = "litellm.bedrock_realtime.pending_session_update" -_COMMITTED_KEY: Final = "litellm.bedrock_realtime.session_committed" _BEDROCK_STREAM_ERROR_STATUS: Final[Mapping[str, int]] = MappingProxyType( { @@ -145,16 +147,14 @@ async def _client_messages(client_ws: RealtimeClientWebSocket, initial_message: return -def _take_pending_session_update( - scope: MutableMapping[str, object], # mutable-ok: the ASGI scope is the per-connection state store -) -> str | None: +def _pending_session_update(scope: Mapping[str, object]) -> str | None: """A fallback attempt on the same websocket replays the session.update the failed attempt never acked.""" - if scope.get(_COMMITTED_KEY) is True: + if scope.get(BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY) is True: raise BedrockError( status_code=409, message="Bedrock realtime session already committed to a provider stream; it cannot be replayed", ) - pending: Final = scope.pop(_PENDING_UPDATE_KEY, None) # rebind-ok: the ASGI scope outlives this attempt + pending: Final = scope.get(BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY) return pending if isinstance(pending, str) else None @@ -175,8 +175,8 @@ async def _ack_session_update( ) -> bool: """Ack the client's session.update once Bedrock accepted the stream; False means the client is gone.""" await bedrock_stream.await_output() - client_ws.scope.pop(_PENDING_UPDATE_KEY, None) - client_ws.scope[_COMMITTED_KEY] = True # rebind-ok: the ASGI scope outlives this attempt + client_ws.scope.pop(BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY, None) + client_ws.scope[BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY] = True # rebind-ok: scope outlives the attempt if logging_obj is None: return True requested_modalities: Final = _CLIENT_MODALITIES_ADAPTER.validate_python( @@ -239,7 +239,7 @@ class BedrockRealtime(BaseAWSLLM): except ImportError: raise ImportError("Missing aws_sdk_bedrock_runtime. Install with: pip install aws-sdk-bedrock-runtime") - pending_session_update: Final = _take_pending_session_update(websocket.scope) + pending_session_update: Final = _pending_session_update(websocket.scope) # Get AWS region if aws_region_name is None: @@ -445,7 +445,9 @@ class BedrockRealtime(BaseAWSLLM): parsed_client_message = _parse_client_message(message) is_session_update = _json_str(parsed_client_message.get("type")) == "session.update" if is_session_update: - client_ws.scope[_PENDING_UPDATE_KEY] = message # rebind-ok: scope outlives the attempt + client_ws.scope[BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY] = ( + message # rebind-ok: scope outlives the attempt + ) transformed_messages = transformation_config.transform_realtime_request( message=message, diff --git a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py index 0439a5090f5..085ffeb5129 100644 --- a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py +++ b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py @@ -217,7 +217,10 @@ def stub_aws_sdk_client(monkeypatch): async def invoke_model_with_bidirectional_stream(self, operation_input): captured["operation_input"] = operation_input if captured.get("streams"): - return captured["streams"].pop(0) + stream = captured["streams"].pop(0) + if isinstance(stream, Exception): + raise stream + return stream return ScriptedBedrockStream(captured.get("scripted_payloads", [])) package = types.ModuleType("aws_sdk_bedrock_runtime") @@ -531,7 +534,8 @@ class TestBedrockRealtimeProviderFailurePropagation: handler = BedrockRealtime() websocket = ConnectedClientWS([self.SESSION_UPDATE]) healthy_stream = ScriptedBedrockStream([]) - stub_aws_sdk_client["streams"] = [UnavailableBedrockStream(), healthy_stream] + eager_failure = ServiceUnavailableException("fault injected before the stream was returned") + stub_aws_sdk_client["streams"] = [UnavailableBedrockStream(), eager_failure, healthy_stream] with pytest.raises(BedrockError) as failure: await handler.async_realtime( @@ -542,6 +546,11 @@ class TestBedrockRealtimeProviderFailurePropagation: assert [json.loads(m)["type"] for m in websocket.sent_to_client] == ["session.created"] assert not websocket.closed, "the proxy route owns the client-facing error event and 1011 close" + with pytest.raises(ServiceUnavailableException): + await handler.async_realtime( + model="amazon.nova-sonic-v1:0", websocket=websocket, logging_obj=FakeLogging(), **self.AWS_PARAMS + ) + await handler.async_realtime( model="amazon.nova-sonic-v1:0", websocket=websocket, logging_obj=FakeLogging(), **self.AWS_PARAMS ) From 771430df7b1ec84830c6f6d66ac1e2fe31497007 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 10:52:16 +0000 Subject: [PATCH 078/112] fix(bedrock/realtime): keep partial spend on cancelled output task and make the committed-session refusal non-retryable Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 1 + litellm/llms/bedrock/realtime/handler.py | 61 +++++----- .../realtime/test_bedrock_realtime_handler.py | 104 ++++++++++++++---- 3 files changed, 112 insertions(+), 54 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 7d86080086a..a1a737153fa 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -314,6 +314,7 @@ WEBSOCKET_CLOSE_REASON_MAX_BYTES: Final = 123 # ASGI websocket scope keys the Bedrock realtime bridge uses to carry state across router fallback attempts BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY: Final = "litellm.bedrock_realtime.pending_session_update" BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY: Final = "litellm.bedrock_realtime.session_committed" +BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY: Final = "litellm.bedrock_realtime.committed_failure" # SSL/TLS cipher configuration for faster handshakes # Strategy: Strongly prefer fast modern ciphers, but allow fallback to commonly supported ones diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 126a7ae5c3e..130dec20a4b 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -10,13 +10,14 @@ import json from collections.abc import AsyncIterator, Mapping, MutableMapping from dataclasses import dataclass from types import MappingProxyType -from typing import Final, Protocol +from typing import Final, NoReturn, Protocol from pydantic import JsonValue, TypeAdapter import litellm from litellm._logging import _redact_string, verbose_proxy_logger from litellm.constants import ( + BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY, BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY, BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY, ) @@ -61,15 +62,6 @@ def _as_bedrock_error(error: BaseException) -> BaseException: return BedrockError(status_code=status_code, message=f"{type(error).__name__}: {error}") -class _BedrockForwardingFailed(Exception): - """The Bedrock output stream failed after ``logged_events`` were already forwarded to the client.""" - - def __init__(self, cause: BaseException, logged_events: tuple[OpenAIRealtimeEvents, ...]) -> None: - super().__init__(str(cause)) - self.cause: Final = cause - self.logged_events: Final = logged_events - - def _json_dict(value: JsonValue) -> dict[str, JsonValue]: return value if isinstance(value, dict) else {} @@ -150,14 +142,26 @@ async def _client_messages(client_ws: RealtimeClientWebSocket, initial_message: def _pending_session_update(scope: Mapping[str, object]) -> str | None: """A fallback attempt on the same websocket replays the session.update the failed attempt never acked.""" if scope.get(BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY) is True: + committed_failure: Final = scope.get(BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY) raise BedrockError( - status_code=409, - message="Bedrock realtime session already committed to a provider stream; it cannot be replayed", + status_code=400, + message=( + "Bedrock realtime session already committed to a provider stream; it cannot be replayed" + + (f". The committed stream failed with: {committed_failure}" if committed_failure else "") + ), ) pending: Final = scope.get(BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY) return pending if isinstance(pending, str) else None +def _raise_provider_failure(scope: MutableMapping[str, object], failure: BaseException) -> NoReturn: + error: Final = _as_bedrock_error(failure) + verbose_proxy_logger.error("Bedrock Realtime: provider stream failed: %s", _redact_string(str(error))) + if scope.get(BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY) is True: + scope[BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY] = _redact_string(str(error)) + raise error from failure + + def _parse_client_message(message: str) -> Mapping[str, JsonValue]: try: return _json_dict(_CLIENT_MESSAGE_ADAPTER.validate_json(message)) @@ -351,10 +355,7 @@ class BedrockRealtime(BaseAWSLLM): "Bedrock Realtime: stream failed after the client disconnected: %s", outcome.provider_failure ) return - verbose_proxy_logger.error( - "Bedrock Realtime: provider stream failed: %s", _redact_string(str(outcome.provider_failure)) - ) - raise _as_bedrock_error(outcome.provider_failure) from outcome.provider_failure + _raise_provider_failure(websocket.scope, outcome.provider_failure) async def _bridge( self, @@ -367,17 +368,13 @@ class BedrockRealtime(BaseAWSLLM): initial_message: str | None, ) -> _BridgeOutcome: """Run both forwarding directions until the client leaves or either side fails.""" + logged: Final[list[OpenAIRealtimeEvents]] = [] # mutable-ok: events forwarded before a failure are still spend - async def collect_logged_events() -> tuple[OpenAIRealtimeEvents, ...]: - logged: Final[list[OpenAIRealtimeEvents]] = [] # mutable-ok: partial events are still logged - try: - async for event in self._forward_bedrock_to_client( - bedrock_stream, websocket, transformation_config, model, logging_obj, session_state - ): - logged.append(event) - except Exception as e: - raise _BedrockForwardingFailed(e, tuple(logged)) from e - return tuple(logged) + async def collect_logged_events() -> None: + async for event in self._forward_bedrock_to_client( + bedrock_stream, websocket, transformation_config, model, logging_obj, session_state + ): + logged.append(event) client_task: Final = asyncio.create_task( self._forward_client_to_bedrock( @@ -395,18 +392,12 @@ class BedrockRealtime(BaseAWSLLM): client_outcome, bedrock_outcome = await asyncio.gather(client_task, bedrock_task, return_exceptions=True) return _BridgeOutcome( - logged_events=( - bedrock_outcome.logged_events - if isinstance(bedrock_outcome, _BedrockForwardingFailed) - else bedrock_outcome - if isinstance(bedrock_outcome, tuple) - else () - ), + logged_events=tuple(logged), provider_failure=( client_outcome if isinstance(client_outcome, Exception) - else bedrock_outcome.cause - if isinstance(bedrock_outcome, _BedrockForwardingFailed) + else bedrock_outcome + if isinstance(bedrock_outcome, Exception) else None ), client_disconnected=client_disconnected, diff --git a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py index 085ffeb5129..2eadaee9e1a 100644 --- a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py +++ b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py @@ -168,6 +168,34 @@ class BreakingBedrockReceiver(ScriptedBedrockReceiver): return await super().receive() +class DrainedThenOpenBedrockReceiver(ScriptedBedrockReceiver): + """Delivers its payloads, flags `drained`, then stays open like a live Nova Sonic turn""" + + def __init__(self, payloads): + super().__init__(payloads) + self.drained = asyncio.Event() + + async def receive(self): + if not self._payloads: + self.drained.set() + await asyncio.Event().wait() + return await super().receive() + + +class ResetOnAudioInputStream(FakeInputStream): + """Accepts session setup, then the provider resets the input side once the first response was delivered""" + + def __init__(self, drained): + super().__init__() + self._drained = drained + + async def send(self, event): + if "audioInput" in json.loads(event.value.bytes_.decode("utf-8")).get("event", {}): + await self._drained.wait() + raise RuntimeError("bedrock input stream reset") + self.sent.append(event) + + class ScriptedBedrockStream: def __init__(self, payloads, receiver_type=ScriptedBedrockReceiver): self.input_stream = FakeInputStream() @@ -560,10 +588,14 @@ class TestBedrockRealtimeProviderFailurePropagation: assert [next(iter(event["event"])) for event in replayed][:2] == ["sessionStart", "promptStart"] assert websocket.closed - @pytest.mark.asyncio - async def test_mid_stream_failure_escapes_keeps_partial_spend_and_blocks_replay( - self, stub_aws_sdk_client, monkeypatch - ): + TEXT_TURN = ( + json.dumps({"event": {"contentStart": {"role": "ASSISTANT", "type": "TEXT"}}}), + json.dumps({"event": {"textOutput": {"content": "Hi"}}}), + json.dumps({"event": {"contentEnd": {"stopReason": "END_TURN"}}}), + ) + + @pytest.fixture + def spend_dispatch(self, monkeypatch): import litellm.llms.bedrock.realtime.handler as handler_module dispatched = {} @@ -577,35 +609,69 @@ class TestBedrockRealtimeProviderFailurePropagation: dispatched["coro"] = coro monkeypatch.setattr(handler_module, "GLOBAL_LOGGING_WORKER", RecordingLoggingWorker()) + dispatched["logging_obj"] = RecordingLogging() + return dispatched + + @pytest.mark.asyncio + async def test_mid_stream_failure_escapes_keeps_partial_spend_and_blocks_replay( + self, stub_aws_sdk_client, spend_dispatch + ): handler = BedrockRealtime() websocket = ConnectedClientWS([self.SESSION_UPDATE]) - stub_aws_sdk_client["streams"] = [ - ScriptedBedrockStream( - [ - json.dumps({"event": {"contentStart": {"role": "ASSISTANT", "type": "TEXT"}}}), - json.dumps({"event": {"textOutput": {"content": "Hi"}}}), - json.dumps({"event": {"contentEnd": {"stopReason": "END_TURN"}}}), - ], - receiver_type=BreakingBedrockReceiver, - ) - ] + stream = ScriptedBedrockStream(self.TEXT_TURN, receiver_type=BreakingBedrockReceiver) + stub_aws_sdk_client["streams"] = [stream] with pytest.raises(BedrockError) as failure: await handler.async_realtime( - model="amazon.nova-sonic-v1:0", websocket=websocket, logging_obj=RecordingLogging(), **self.AWS_PARAMS + model="amazon.nova-sonic-v1:0", + websocket=websocket, + logging_obj=spend_dispatch["logging_obj"], + **self.AWS_PARAMS, ) assert failure.value.status_code == 424 - await dispatched["coro"] - assert [event["type"] for event in dispatched["events"]] == ["response.done"] + await spend_dispatch["coro"] + assert [event["type"] for event in spend_dispatch["events"]] == ["response.done"] assert "response.done" in [json.loads(m)["type"] for m in websocket.sent_to_client] + flushed = [json.loads(chunk.value.bytes_.decode("utf-8")) for chunk in stream.input_stream.sent] + assert [next(iter(event["event"])) for event in flushed][-2:] == ["promptEnd", "sessionEnd"] + assert stream.input_stream.closed with pytest.raises(BedrockError) as replay: await handler.async_realtime( - model="amazon.nova-sonic-v1:0", websocket=websocket, logging_obj=RecordingLogging(), **self.AWS_PARAMS + model="amazon.nova-sonic-v1:0", + websocket=websocket, + logging_obj=spend_dispatch["logging_obj"], + **self.AWS_PARAMS, ) - assert replay.value.status_code == 409, "a committed session must not be silently restarted on a fallback" + assert replay.value.status_code == 400, "a committed session must not be silently restarted on a fallback" + assert not litellm._should_retry(replay.value.status_code), "the router must not retry the replay refusal" + assert "Nova Sonic stream broke" in replay.value.message, "the router surfaces the last attempt's error" + + @pytest.mark.asyncio + async def test_input_side_failure_keeps_spend_for_responses_already_delivered( + self, stub_aws_sdk_client, spend_dispatch + ): + receiver = DrainedThenOpenBedrockReceiver(self.TEXT_TURN) + stream = ScriptedBedrockStream(self.TEXT_TURN, receiver_type=lambda _payloads: receiver) + stream.input_stream = ResetOnAudioInputStream(receiver.drained) + stub_aws_sdk_client["streams"] = [stream] + websocket = ConnectedClientWS( + [self.SESSION_UPDATE, json.dumps({"type": "input_audio_buffer.append", "audio": "AAAA"})] + ) + + with pytest.raises(RuntimeError, match="bedrock input stream reset"): + await BedrockRealtime().async_realtime( + model="amazon.nova-sonic-v1:0", + websocket=websocket, + logging_obj=spend_dispatch["logging_obj"], + **self.AWS_PARAMS, + ) + + assert "response.done" in [json.loads(m)["type"] for m in websocket.sent_to_client] + await spend_dispatch["coro"] + assert [event["type"] for event in spend_dispatch["events"]] == ["response.done"] @pytest.mark.asyncio async def test_stream_failure_after_client_disconnect_is_not_a_provider_failure(self, stub_aws_sdk_client): From db378d896349b7ee0ca02cf48d0cb266f4b95736 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 14 Sep 2026 12:08:31 -0700 Subject: [PATCH 079/112] fix(harness): preserve unmapped calls in execution traces --- tests/rust-python-harness/AGENTS.md | 2 +- .../shared/tracing/steps.py | 4 +- .../shared/tracing/test_steps.py | 17 ++++ .../strategies/trace_parity/AGENTS.md | 2 +- .../trace_parity/gateway/execution.py | 8 +- .../strategies/trace_parity/models.py | 2 + .../strategies/trace_parity/sdk/execution.py | 10 +-- .../strategies/trace_parity/sdk/ocr/case.py | 46 +++++++++++ .../strategies/trace_parity/test_runner.py | 80 ++++++++++++++++++- 9 files changed, 156 insertions(+), 15 deletions(-) diff --git a/tests/rust-python-harness/AGENTS.md b/tests/rust-python-harness/AGENTS.md index ec973035785..71e17541fd2 100644 --- a/tests/rust-python-harness/AGENTS.md +++ b/tests/rust-python-harness/AGENTS.md @@ -63,7 +63,7 @@ tests/rust-python-harness/ - Examples: `run e2e_parity --surface sdk --function ocr`, `run unit_tests_parity --function ocr --pytest-arg=-x`, or `run all --function ocr` - `cli/catalog.py` discovers strategies, validates their Python definitions, and orders them; `cli/__init__.py` builds the Click command tree; `cli/commands.py` runs selected cases - `e2e_parity/` compares SDK objects, exceptions, callbacks, and streams, or gateway HTTP responses -- `trace_parity/` prints filtered Python and Rust execution traces without comparing them; before running it rebuilds the native bridge with the `trace-parity` feature whenever `litellm-rust` sources are newer than the installed extension (`shared/native_build.py`) +- `trace_parity/` prints every collected Python call under `litellm/` and every Rust span without comparing them; mappings only filter the separate unit-test mapping strategy. Before running it rebuilds the native bridge with the `trace-parity` feature whenever `litellm-rust` sources are newer than the installed extension (`shared/native_build.py`) - E2E and trace strategies load their registered module cases and run surface-specific execution from their folders - `unit_tests_mapping/contracts.py` owns typed harness-side mapping contracts, per-function contracts live below `cases/`, and `mappings.py` exports the registry; live test discovery derives unmapped Python and Rust-only tests without an exhaustive manifest - `unit_tests_mapping/runner.py` validates confirmed mappings against the live Python and Rust inventories and attaches the derived status report diff --git a/tests/rust-python-harness/shared/tracing/steps.py b/tests/rust-python-harness/shared/tracing/steps.py index 152f753a536..492ffab64e5 100644 --- a/tests/rust-python-harness/shared/tracing/steps.py +++ b/tests/rust-python-harness/shared/tracing/steps.py @@ -76,7 +76,7 @@ def _span_for(engine: Engine, function: str, mappings: Sequence[TraceMapping]) - def pipeline_projection( - engine: Engine, events: Sequence[FunctionTraceEvent], mappings: Sequence[TraceMapping] + engine: Engine, events: Sequence[FunctionTraceEvent], mappings: Sequence[TraceMapping] | None = None ) -> PipelineProjection: raw_parents: dict[int, int | None] = {} projected_ids: set[int] = set() @@ -88,7 +88,7 @@ def pipeline_projection( if event.parent_id is not None and event.parent_id not in raw_parents: raise ValueError(f"trace event {event.id} references unknown or later parent {event.parent_id}") raw_parents[event.id] = event.parent_id - span = _span_for(engine, event.function, mappings) + span = event.function if mappings is None else _span_for(engine, event.function, mappings) if span is None: unmatched += 1 continue diff --git a/tests/rust-python-harness/shared/tracing/test_steps.py b/tests/rust-python-harness/shared/tracing/test_steps.py index 1b2f02a4ed7..ee5e0bafd28 100644 --- a/tests/rust-python-harness/shared/tracing/test_steps.py +++ b/tests/rust-python-harness/shared/tracing/test_steps.py @@ -40,6 +40,23 @@ def test_python_projection_collapses_unmapped_parents_and_counts_noise() -> None ] +@pytest.mark.parametrize("engine", ("python", "rust")) +def test_projection_without_mappings_keeps_every_call_and_parent(engine: Engine) -> None: + events: Final = ( + event(0, "module.py:1 entry"), + event(1, "module.py:2 internal_helper", 0), + event(2, "module.py:3 nested", 1), + event(3, "module.py:2 internal_helper", 0), + ) + + projection: Final = pipeline_projection(engine, events) + + assert projection.unmatched == 0 + assert tuple((step.id, step.parent_id, step.span, step.raw) for step in projection.steps) == tuple( + (item.id, item.parent_id, item.function, item.raw) for item in events + ) + + def test_rust_projection_keeps_unknown_spans() -> None: projection: Final = pipeline_projection("rust", (event(0, "route"), event(1, "new_span", 0)), MAPPINGS) assert [(step.span, step.parent_id) for step in projection.steps] == [("route", None), ("new_span", 0)] diff --git a/tests/rust-python-harness/strategies/trace_parity/AGENTS.md b/tests/rust-python-harness/strategies/trace_parity/AGENTS.md index 5db17974d07..19861aea2b4 100644 --- a/tests/rust-python-harness/strategies/trace_parity/AGENTS.md +++ b/tests/rust-python-harness/strategies/trace_parity/AGENTS.md @@ -1 +1 @@ -Prints filtered Python profiler frames and feature-gated Rust spans from live traces against a replayed provider response. The two traces are independent and are not compared. +Prints every collected Python call under litellm/ and every feature-gated Rust span from live traces against replayed HTTP responses. The two traces are independent and are not compared. API-key and Vertex credentials scenarios exercise separate authentication paths; credentials scenarios replay the token exchange locally. diff --git a/tests/rust-python-harness/strategies/trace_parity/gateway/execution.py b/tests/rust-python-harness/strategies/trace_parity/gateway/execution.py index 88130dab808..5f24760dfe3 100644 --- a/tests/rust-python-harness/strategies/trace_parity/gateway/execution.py +++ b/tests/rust-python-harness/strategies/trace_parity/gateway/execution.py @@ -160,13 +160,11 @@ def _collect( def _projections( python_events: tuple[FunctionTraceEvent, ...], rust_events: tuple[FunctionTraceEvent, ...], - scenario: TraceScenario, ) -> tuple[PipelineProjection, PipelineProjection, str | None]: - mappings: Final = scenario.mappings try: return ( - pipeline_projection("python", python_events, mappings), - pipeline_projection("rust", rust_events, mappings), + pipeline_projection("python", python_events), + pipeline_projection("rust", rust_events), None, ) except ValueError as error: @@ -184,7 +182,7 @@ def execute_gateway_trace( rust_error: Final = None if isinstance(rust_trace, tuple) else f"rust: {rust_trace.message}" python_events: Final = python_trace if isinstance(python_trace, tuple) else () rust_events: Final = rust_trace if isinstance(rust_trace, tuple) else () - python, rust, projection_error = _projections(python_events, rust_events, scenario) + python, rust, projection_error = _projections(python_events, rust_events) python_error: Final = projection_error or collection_python_error return TraceArtifact.from_traces( engine=engine, diff --git a/tests/rust-python-harness/strategies/trace_parity/models.py b/tests/rust-python-harness/strategies/trace_parity/models.py index be7375dc604..fd4585d85c1 100644 --- a/tests/rust-python-harness/strategies/trace_parity/models.py +++ b/tests/rust-python-harness/strategies/trace_parity/models.py @@ -18,6 +18,7 @@ class RouteFixture: provider_responses: tuple[RecordedResponse, ...] expected_failure: bool = False consume_stream: bool = False + environment: tuple[tuple[str, str], ...] = () def derive( self, @@ -32,6 +33,7 @@ class RouteFixture: provider_responses=self.provider_responses if provider_responses is None else provider_responses, expected_failure=self.expected_failure if expected_failure is None else expected_failure, consume_stream=self.consume_stream if consume_stream is None else consume_stream, + environment=self.environment, ) def with_body(self, **updates: object) -> RouteFixture: diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py b/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py index f8a674451b8..2ef8a0584bc 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py @@ -120,21 +120,22 @@ def collect_trace( provider.enqueue_response(response) fixture: Final = RouteFixture( kwargs={ - **base_fixture.kwargs, "api_key": "test-key", + **base_fixture.kwargs, "api_base": provider.url, **({"timeout_seconds": 5} if engine == "rust" else {"timeout": 5}), }, provider_responses=base_fixture.provider_responses, expected_failure=base_fixture.expected_failure, consume_stream=base_fixture.consume_stream, + environment=base_fixture.environment, ) environment: Final = ( patch.dict(os.environ, {"LITELLM_RUST": "1" if python_rust_enabled else "0"}) if engine == "python" else nullcontext() ) - with environment: + with environment, patch.dict(os.environ, fixture.environment): collected: Final = _collect(function, fixture, engine, asynchronous=asynchronous) provider.take_requests(len(fixture.provider_responses)) except Exception as error: @@ -160,7 +161,6 @@ def execute_trace( surface: Surface, engine: TraceEngine = "both", ) -> TraceArtifact: - mappings: Final = scenario.mappings scenario_route: Final = RouteSpec( route=route.route, python_entrypoints=route.python_entrypoints, @@ -185,8 +185,8 @@ def execute_trace( python_events: Final = python_trace if isinstance(python_trace, tuple) else () rust_events: Final = rust_trace if isinstance(rust_trace, tuple) else () try: - python: Final = pipeline_projection("python", python_events, mappings) - rust: Final = pipeline_projection("rust", rust_events, mappings) + python: Final = pipeline_projection("python", python_events) + rust: Final = pipeline_projection("rust", rust_events) except ValueError as error: return TraceArtifact.from_traces( engine=engine, diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py b/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py index b2c0e4c7acc..58b16c68383 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py @@ -204,6 +204,40 @@ def _vertex_deepseek_fixture(engine: Engine, _base_url: str) -> RouteFixture: ) +def _vertex_deepseek_credentials_fixture(engine: Engine, base_url: str) -> RouteFixture: + from cryptography.hazmat.primitives import serialization + from cryptography.hazmat.primitives.asymmetric import rsa + + fixture: Final = _vertex_deepseek_fixture(engine, base_url) + private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + credentials: Final = json.dumps( + { + "type": "service_account", + "project_id": "trace-project", + "private_key_id": "trace-key", + "private_key": private_key.private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.PKCS8, + serialization.NoEncryption(), + ).decode(), + "client_email": "trace@trace-project.iam.gserviceaccount.com", + "token_uri": f"{base_url}/token", + } + ) + return RouteFixture( + kwargs={**fixture.kwargs, "api_key": None}, + environment=(("VERTEXAI_CREDENTIALS", credentials), ("VERTEX_AI_API_KEY", "")), + provider_responses=( + RecordedHttpResponse.from_bytes( + 200, + (HttpHeader(name="content-type", value="application/json"),), + b'{"access_token":"trace-token","token_type":"Bearer","expires_in":3600}', + ), + *fixture.provider_responses, + ), + ) + + def _cohere_fixture(engine: Engine, _base_url: str) -> RouteFixture: return RouteFixture( kwargs={ @@ -441,6 +475,18 @@ TRACE_SUITE: Final = TraceSuite( mappings=(*DEEPSEEK_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), asynchronous=True, ), + TraceScenario( + name="sync-vertex-deepseek-credentials", + fixture=_vertex_deepseek_credentials_fixture, + mappings=DEEPSEEK_SYNC_MAPPINGS, + asynchronous=False, + ), + TraceScenario( + name="async-vertex-deepseek-credentials", + fixture=_vertex_deepseek_credentials_fixture, + mappings=DEEPSEEK_ASYNC_MAPPINGS, + asynchronous=True, + ), TraceScenario( name="async-cohere", fixture=_cohere_fixture, diff --git a/tests/rust-python-harness/strategies/trace_parity/test_runner.py b/tests/rust-python-harness/strategies/trace_parity/test_runner.py index 9dfd0d51c89..b940221bb7a 100644 --- a/tests/rust-python-harness/strategies/trace_parity/test_runner.py +++ b/tests/rust-python-harness/strategies/trace_parity/test_runner.py @@ -13,7 +13,7 @@ import litellm from ...shared.reporting.models import Coverage, HarnessCase, HarnessRun, RunStatus, SdkFunction, Surface from ...shared.reporting.strategy import ModuleCaseSpec from ...shared.tracing.profiler import FunctionTraceEvent -from ...shared.tracing.steps import Engine, PipelineStep +from ...shared.tracing.steps import Engine, PipelineStep, mapping from .models import GatewayRouteSpec, RouteFixture, RouteSpec, TraceScenario, TraceSuite from .reporting import TraceArtifact from .runner import run_trace_cases, run_trace_scenario, runner_selection, scenario_nodeids, validate_trace_suite @@ -124,6 +124,84 @@ def test_expected_provider_failure_omits_feedback_banner( assert litellm.suppress_debug_info is False +@pytest.mark.parametrize("asynchronous", (False, True)) +def test_vertex_trace_keeps_unmapped_helpers_and_parents(asynchronous: bool) -> None: + loaded: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.sdk.ocr.case") + suite: Final = cast(TraceSuite, loaded.TRACE_SUITE) + name: Final = f"{'async' if asynchronous else 'sync'}-vertex-deepseek" + scenario: Final = next(item for item in suite.scenarios if item.name == name) + assert isinstance(suite.route, RouteSpec) + + trace: Final = execute_trace(suite.route, scenario, "sdk", engine="python") + + assert trace.python_error is None + url: Final = next( + event for event in trace.python if event.raw.endswith(" VertexAIDeepSeekOCRConfig.get_complete_url") + ) + project: Final = next( + event for event in trace.python if event.raw.endswith(" VertexBase.safe_get_vertex_ai_project") + ) + location: Final = next( + event for event in trace.python if event.raw.endswith(" VertexBase.safe_get_vertex_ai_location") + ) + assert project.parent_id == location.parent_id == url.id + assert not any(event.raw.endswith(" VertexBase.get_access_token") for event in trace.python) + + +@pytest.mark.parametrize("asynchronous", (False, True)) +def test_vertex_credentials_trace_runs_real_auth_helpers(asynchronous: bool, monkeypatch: pytest.MonkeyPatch) -> None: + loaded: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.sdk.ocr.case") + suite: Final = cast(TraceSuite, loaded.TRACE_SUITE) + name: Final = f"{'async' if asynchronous else 'sync'}-vertex-deepseek-credentials" + scenario: Final = next(item for item in suite.scenarios if item.name == name) + monkeypatch.setenv("VERTEXAI_CREDENTIALS", "original-credentials") + monkeypatch.setenv("VERTEX_AI_API_KEY", "original-api-key") + assert isinstance(suite.route, RouteSpec) + + trace: Final = execute_trace(suite.route, scenario, "sdk", engine="python") + + assert trace.python_error is None + validate: Final = next( + event for event in trace.python if event.raw.endswith(" VertexAIDeepSeekOCRConfig.validate_environment") + ) + helpers: Final = ( + "VertexBase.safe_get_vertex_ai_project", + "VertexBase.safe_get_vertex_ai_credentials", + "VertexBase.get_access_token", + ) + assert tuple(event.raw.split(" ", 1)[1] for event in trace.python if event.parent_id == validate.id) == helpers + token: Final = next(event for event in trace.python if event.raw.endswith(" VertexBase.get_access_token")) + load: Final = next(event for event in trace.python if event.raw.endswith(" VertexBase.load_auth")) + refresh: Final = next(event for event in trace.python if event.raw.endswith(" VertexBase.refresh_auth")) + assert load.parent_id == token.id + assert refresh.parent_id == load.id + assert os.environ["VERTEXAI_CREDENTIALS"] == "original-credentials" + assert os.environ["VERTEX_AI_API_KEY"] == "original-api-key" + + +def test_gateway_trace_keeps_calls_outside_scenario_mappings(monkeypatch: pytest.MonkeyPatch) -> None: + execution: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.gateway.execution") + events: Final = ( + FunctionTraceEvent(0, None, "route.py:1 entry"), + FunctionTraceEvent(1, 0, "auth.py:2 authenticate"), + FunctionTraceEvent(2, 1, "auth.py:3 credentials"), + ) + scenario: Final = TraceScenario( + "async-gateway", + _fixture, + (mapping(rust_span="entry", python_frame=r" entry$"),), + asynchronous=True, + ) + monkeypatch.setattr(execution, "_collect", lambda *_args: events) + + trace: Final = execution.execute_gateway_trace(GatewayRouteSpec("messages"), scenario, engine="python") + + assert trace.python_error is None + assert tuple((event.id, event.parent_id, event.raw) for event in trace.python) == tuple( + (event.id, event.parent_id, event.raw) for event in events + ) + + def test_scenario_validation_rejects_duplicate_and_unsafe_names() -> None: route: Final = RouteSpec("ocr", ("ocr", "aocr"), ("ocr", "aocr"), _fixture) duplicate: Final = TraceSuite( From 5bcc23014e8e64ea07a7e9259e06ef86b6e051d8 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 19:12:35 +0000 Subject: [PATCH 080/112] refactor(auth): drop the membership fetch docstring Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/auth_checks.py | 1 - 1 file changed, 1 deletion(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index a8f9cc77577..fa85400f4e0 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2197,7 +2197,6 @@ async def _fetch_team_membership_from_db( parent_otel_span: Span | None = None, proxy_logging_obj: ProxyLogging | None = None, ) -> LiteLLM_TeamMembership | None: - """Prisma read + cache populate. Decorated so cache hits on ``get_team_membership`` are not postgres spans.""" _ = parent_otel_span, proxy_logging_obj response: Final = await _dictable_table(TeamMembershipRepository(prisma_client)).find_unique( where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, From 27981c7d20c14a54c0604037bd2312bf887c86f8 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 14 Sep 2026 12:17:15 -0700 Subject: [PATCH 081/112] fix: preserve rust setting in trace parity --- .../strategies/trace_parity/models.py | 1 - .../strategies/trace_parity/sdk/execution.py | 13 ++----------- .../strategies/trace_parity/sdk/ocr/case.py | 2 -- .../strategies/trace_parity/test_runner.py | 7 ++++--- 4 files changed, 6 insertions(+), 17 deletions(-) diff --git a/tests/rust-python-harness/strategies/trace_parity/models.py b/tests/rust-python-harness/strategies/trace_parity/models.py index fd4585d85c1..3701fb41e14 100644 --- a/tests/rust-python-harness/strategies/trace_parity/models.py +++ b/tests/rust-python-harness/strategies/trace_parity/models.py @@ -66,7 +66,6 @@ class TraceScenario: fixture: Callable[[Engine, str], RouteFixture] mappings: tuple[TraceMapping, ...] asynchronous: bool - python_rust_enabled: bool = False @dataclass(frozen=True, slots=True) diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py b/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py index 2ef8a0584bc..ffdbe842e85 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py @@ -3,7 +3,6 @@ from __future__ import annotations import asyncio import os from collections.abc import AsyncIterable, Awaitable, Iterable -from contextlib import nullcontext from dataclasses import dataclass from pathlib import Path from typing import Final, Protocol, cast @@ -107,9 +106,7 @@ def _collect( return _CollectedTrace(tuple(profiler.events), error) -def collect_trace( - spec: RouteSpec, engine: Engine, *, asynchronous: bool, python_rust_enabled: bool = False -) -> tuple[FunctionTraceEvent, ...] | TraceExecutionFailure: +def collect_trace(spec: RouteSpec, engine: Engine, *, asynchronous: bool) -> tuple[FunctionTraceEvent, ...] | TraceExecutionFailure: function: Final = _entrypoint(spec, engine, asynchronous=asynchronous) if isinstance(function, TraceExecutionFailure): return function @@ -130,12 +127,7 @@ def collect_trace( consume_stream=base_fixture.consume_stream, environment=base_fixture.environment, ) - environment: Final = ( - patch.dict(os.environ, {"LITELLM_RUST": "1" if python_rust_enabled else "0"}) - if engine == "python" - else nullcontext() - ) - with environment, patch.dict(os.environ, fixture.environment): + with patch.dict(os.environ, fixture.environment): collected: Final = _collect(function, fixture, engine, asynchronous=asynchronous) provider.take_requests(len(fixture.provider_responses)) except Exception as error: @@ -172,7 +164,6 @@ def execute_trace( scenario_route, "python", asynchronous=scenario.asynchronous, - python_rust_enabled=scenario.python_rust_enabled, ) if engine != "rust" else () diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py b/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py index 58b16c68383..bb21e8ab0c5 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py @@ -498,14 +498,12 @@ TRACE_SUITE: Final = TraceSuite( fixture=_mistral_fixture, mappings=PUBLIC_RUST_DISPATCH_MAPPINGS, asynchronous=False, - python_rust_enabled=True, ), TraceScenario( name="async-public-rust-dispatch", fixture=_mistral_fixture, mappings=PUBLIC_RUST_DISPATCH_MAPPINGS, asynchronous=True, - python_rust_enabled=True, ), ), ) diff --git a/tests/rust-python-harness/strategies/trace_parity/test_runner.py b/tests/rust-python-harness/strategies/trace_parity/test_runner.py index b940221bb7a..57364057788 100644 --- a/tests/rust-python-harness/strategies/trace_parity/test_runner.py +++ b/tests/rust-python-harness/strategies/trace_parity/test_runner.py @@ -80,7 +80,7 @@ def test_python_engine_skips_native_bridge(monkeypatch: pytest.MonkeyPatch, tmp_ assert selected == [(frozenset({"mistral"}), "python")] -def test_python_trace_controls_native_ocr_dispatch(monkeypatch: pytest.MonkeyPatch) -> None: +def test_python_trace_preserves_native_ocr_dispatch_setting(monkeypatch: pytest.MonkeyPatch) -> None: execution: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.sdk.execution") route: Final = RouteSpec("ocr", ("ocr", "aocr"), ("ocr", "aocr"), _fixture) observed: list[str | None] = [] @@ -98,11 +98,12 @@ def test_python_trace_controls_native_ocr_dispatch(monkeypatch: pytest.MonkeyPat error=None, ) - monkeypatch.setenv("LITELLM_RUST", "1") monkeypatch.setattr(execution, "_collect", collect) + monkeypatch.setenv("LITELLM_RUST", "0") collect_trace(route, "python", asynchronous=False) - collect_trace(route, "python", asynchronous=True, python_rust_enabled=True) + monkeypatch.setenv("LITELLM_RUST", "1") + collect_trace(route, "python", asynchronous=True) assert observed == ["0", "1"] assert os.environ["LITELLM_RUST"] == "1" From 9ff0dea1396e48c302fe203c44f4643f228dc50f Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 19:25:03 +0000 Subject: [PATCH 082/112] fix(proxy): define credentials hint helper ahead of proxy_server imports Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/auth_utils.py | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 0f4160d840b..d882141fa6f 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -552,6 +552,18 @@ def _coerce_metadata_to_dict(value: Any) -> dict[str, Any] | None: return None +def should_hide_default_credentials_hint(general_settings: Mapping[str, object]) -> bool: + """ + Whether login pages hide the "admin / MASTER_KEY" hint: explicit opt-in, or a + non-empty UI_PASSWORD, which makes that hint wrong. UI_USERNAME alone keeps it. + """ + return ( + os.getenv("LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT", "false").lower() == "true" + or general_settings.get("hide_default_credentials_hint", False) is True + or bool(os.getenv("UI_PASSWORD")) + ) + + async def pre_db_read_auth_checks( request: Request, request_data: dict, @@ -1446,18 +1458,6 @@ def has_user_setup_sso() -> bool: ) -def should_hide_default_credentials_hint(general_settings: Mapping[str, object]) -> bool: - """ - Whether login pages hide the "admin / MASTER_KEY" hint: explicit opt-in, or a - non-empty UI_PASSWORD, which makes that hint wrong. UI_USERNAME alone keeps it. - """ - return ( - os.getenv("LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT", "false").lower() == "true" - or general_settings.get("hide_default_credentials_hint", False) is True - or bool(os.getenv("UI_PASSWORD")) - ) - - def _is_google_ready() -> bool: return bool(os.getenv("GOOGLE_CLIENT_ID")) and bool(os.getenv("GOOGLE_CLIENT_SECRET")) From ff023906855c96cb1d2a65e4fc02cd9882302761 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 19:26:14 +0000 Subject: [PATCH 083/112] test(auth): drop docstrings from team membership tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/proxy/auth/test_auth_checks.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 04e41608b26..b843916debf 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -6456,7 +6456,6 @@ async def test_get_team_membership_reads_sentinel_as_no_membership_not_a_model() @pytest.mark.asyncio async def test_get_team_membership_coalesces_parallel_db_fetches(): - """Concurrent misses for the same member must share one Prisma round-trip.""" from litellm.proxy.auth.auth_checks import get_team_membership started = asyncio.Event() @@ -6496,7 +6495,6 @@ async def test_get_team_membership_coalesces_parallel_db_fetches(): @pytest.mark.asyncio async def test_common_checks_calls_get_team_membership_once_per_request(): - """Model-access, attribution, and member-budget must reuse one membership load.""" from fastapi import Request from litellm.proxy.auth.auth_checks import common_checks @@ -6543,7 +6541,6 @@ async def test_common_checks_calls_get_team_membership_once_per_request(): @pytest.mark.asyncio async def test_common_checks_skips_membership_load_when_no_check_reads_it(): - """A management route has no model and no budget gate, so the membership row is never loaded.""" from fastapi import Request from litellm.proxy.auth.auth_checks import common_checks @@ -6583,7 +6580,6 @@ async def test_common_checks_skips_membership_load_when_no_check_reads_it(): @pytest.mark.asyncio async def test_get_team_membership_db_error_returns_none_and_retries_next_call(): - """A Prisma failure reads as no membership, caches nothing, and the next call hits the DB again.""" from litellm.proxy.auth.auth_checks import get_team_membership from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key @@ -6620,7 +6616,6 @@ async def test_get_team_membership_db_error_returns_none_and_retries_next_call() @pytest.mark.asyncio async def test_get_team_membership_string_prisma_client_returns_none(): - """Unit tests stub prisma_client as a string; the lookup fails and reads as no membership.""" from litellm.proxy.auth.auth_checks import get_team_membership result = await get_team_membership( @@ -6634,7 +6629,6 @@ async def test_get_team_membership_string_prisma_client_returns_none(): @pytest.mark.asyncio async def test_get_team_membership_waiter_cancel_does_not_cancel_shared_load(): - """Cancelling one coalesced waiter must not cancel the shared Prisma load.""" from litellm.proxy.auth.auth_checks import get_team_membership started = asyncio.Event() From ebcd9bcb18a3cfae179df42ae0464384a7aa31e0 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 19:29:11 +0000 Subject: [PATCH 084/112] fix(proxy): release max_parallel_requests slot when a realtime session ends without LLM callbacks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 2 + tests/test_litellm/proxy/test_proxy_server.py | 40 +++++++++++++++++++ 2 files changed, 42 insertions(+) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d9f8e04ebda..ed3d43800f7 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11912,6 +11912,7 @@ async def _reject_realtime_session( await websocket.close(code=code, reason=reason) finally: await _release_realtime_budget_reservation(user_api_key_dict) + await proxy_logging_obj._arelease_max_parallel_requests_on_disconnect(user_api_key_dict) # pyright: ignore[reportPrivateUsage] # same release idiom the HTTP disconnect path uses @app.websocket("/openai/v1/realtime") @@ -12050,6 +12051,7 @@ async def realtime_websocket_endpoint( if not litellm_logging_obj.model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY): await _release_realtime_budget_reservation(user_api_key_dict) + await proxy_logging_obj._arelease_max_parallel_requests_on_disconnect(user_api_key_dict) # pyright: ignore[reportPrivateUsage] # same release idiom the HTTP disconnect path uses ###################################################################### diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 03a24ec5e98..7b6a838aa3e 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -10068,6 +10068,46 @@ async def test_successful_realtime_session_leaves_the_reservation_for_the_cost_c assert reservation["finalized"] is False +@pytest.mark.asyncio +@pytest.mark.parametrize("phase_one_exit", [None, "pre_call"]) +async def test_realtime_session_ending_without_llm_callbacks_releases_the_max_parallel_slot( + phase_one_exit: str | None, +): + """The rate limiter acquires the key's max_parallel_requests slot in pre-call and + only frees it from the LLM success/failure callbacks. A realtime session that ends + without either callback (Bedrock closes without usage events, or a later pre-call + hook rejects the session) has to be released by the route itself, or the slot stays + occupied until its TTL and the key's next session is refused with a 429.""" + from litellm.caching.caching import DualCache + from litellm.proxy import proxy_server as ps + from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + RequestRateLimiterStash, + _PROXY_MaxParallelRequestsHandler_v3, + _request_stash, + ) + from litellm.proxy.utils import InternalUsageCache + + counter_key: Final = "{api_key:hashed-token}:max_parallel_requests" + dual_cache: Final = DualCache() + await dual_cache.async_set_cache(key=counter_key, value={"slot-1": 1.0, "slot-2": 2.0}, local_only=True) + limiter: Final = _PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(dual_cache)) + stash: Final = RequestRateLimiterStash(parallel_slot={"slot_id": "slot-1", "counter_keys": [counter_key]}) + reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} + + stash_token: Final = _request_stash.set(stash) + try: + hooks: Final = patch.dict(ps.proxy_logging_obj.proxy_hook_mapping, {"parallel_request_limiter": limiter}) # test-quality-ok: registers a real limiter on the module-global hook map the route reads; assertion observes its counter + with hooks: + await _lit6973_drive_realtime_session( + reservation, backend_logged_success=False, phase_one_exit=phase_one_exit + ) + finally: + _request_stash.reset(stash_token) + + assert await dual_cache.async_get_cache(key=counter_key, local_only=True) == {"slot-2": 2.0} + assert stash.parallel_slot is None + + @pytest.mark.asyncio async def test_release_or_invalidate_falls_back_to_invalidating_the_counters(): """If releasing the reservation itself fails (e.g. the counter store is down), From f9d0cd64097665f0ab4c872cdc3fe7ffae75d507 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 19:29:30 +0000 Subject: [PATCH 085/112] chore(proxy): drop explanatory docstrings from credentials hint helper and tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/auth_utils.py | 4 ---- .../proxy/discovery_endpoints/test_ui_discovery_endpoints.py | 2 -- tests/test_litellm/proxy/management_endpoints/test_ui_sso.py | 2 -- .../test_litellm/proxy/proxy_server/test_routes_login_sso.py | 1 - 4 files changed, 9 deletions(-) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index d882141fa6f..128b789ee75 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -553,10 +553,6 @@ def _coerce_metadata_to_dict(value: Any) -> dict[str, Any] | None: def should_hide_default_credentials_hint(general_settings: Mapping[str, object]) -> bool: - """ - Whether login pages hide the "admin / MASTER_KEY" hint: explicit opt-in, or a - non-empty UI_PASSWORD, which makes that hint wrong. UI_USERNAME alone keeps it. - """ return ( os.getenv("LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT", "false").lower() == "true" or general_settings.get("hide_default_credentials_hint", False) is True diff --git a/tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py b/tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py index 5e3d9c0e43b..64a2eb69325 100644 --- a/tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py +++ b/tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py @@ -350,7 +350,6 @@ def test_ui_discovery_endpoints_hide_default_credentials_hint_default_false(): def test_ui_discovery_endpoints_hide_default_credentials_hint_when_ui_password_set(): - """A non-empty UI_PASSWORD hides the card: 'admin / MASTER_KEY' is no longer the login.""" app = FastAPI() app.include_router(router) client = TestClient(app) @@ -372,7 +371,6 @@ def test_ui_discovery_endpoints_hide_default_credentials_hint_when_ui_password_s ], ) def test_ui_discovery_endpoints_keeps_default_credentials_hint_without_real_ui_password(env_overrides): - """Only a non-empty UI_PASSWORD counts as custom credentials; the hint stays accurate otherwise.""" app = FastAPI() app.include_router(router) client = TestClient(app) diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 2e3c2ceffa3..1230c548281 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -8389,8 +8389,6 @@ async def test_legacy_login_page_hides_credentials_hint_via_general_settings(): @pytest.mark.asyncio async def test_legacy_login_page_hides_credentials_hint_when_ui_password_set(): - """Regression: the legacy page shares the rule with the discovery endpoint, so a non-empty - UI_PASSWORD hides the now-inaccurate 'admin / MASTER_KEY' hint here too.""" response = await _render_legacy_login_page( env_overrides={"UI_PASSWORD": "s3cret-pass"}, general_settings={}, diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_login_sso.py b/tests/test_litellm/proxy/proxy_server/test_routes_login_sso.py index 1b4f9bc1cfe..79c23b11f3e 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_login_sso.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_login_sso.py @@ -106,7 +106,6 @@ def test_fallback_login_shows_credentials_hint_by_default(client, monkeypatch): def test_fallback_login_hides_credentials_hint_when_ui_password_set(client, monkeypatch): - """Regression: a non-empty UI_PASSWORD means 'admin / MASTER_KEY' is wrong, so the hint must go.""" monkeypatch.setenv("UI_PASSWORD", "s3cret-pass") monkeypatch.delenv("LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT", raising=False) response = client.get("/fallback/login") From c4354c248a78efec9842c59c1dee13c7e556310f Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 19:38:14 +0000 Subject: [PATCH 086/112] fix(redis): route the per-TTL pipeline write timeout log through the shared throttle Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/caching/redis_cache.py | 7 +++++-- tests/test_litellm/caching/test_redis_cache.py | 4 ++++ 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 3c3539eb8a9..b4b2b1a334c 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -1231,8 +1231,11 @@ class RedisCache(BaseCache): end_time=time.time(), ) ) - verbose_logger.error( - "LiteLLM Redis Caching: async_set_cache_pipeline_with_ttls() - Got exception from REDIS %s", str(e) + log_redis_failure( + verbose_logger, + logging.ERROR, + "LiteLLM Redis Caching: async_set_cache_pipeline_with_ttls() - Got exception from REDIS", + e, ) _record_swallowed_redis_failure(self._circuit_breaker, e) diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index 9fdddd3c085..19638c60b4b 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -1260,6 +1260,10 @@ def test_timeouts_during_a_blip_log_once_per_interval_not_once_per_call(sync_bat "call_method", [ pytest.param(lambda c: c.async_set_cache_pipeline([("lit7520", "v")]), id="async_set_cache_pipeline"), + pytest.param( + lambda c: c.async_set_cache_pipeline_with_ttls([("lit7520", "v", 60.0)]), + id="async_set_cache_pipeline_with_ttls", + ), pytest.param(lambda c: c.async_set_cache_sadd("lit7520", ["v"], ttl=None), id="async_set_cache_sadd"), pytest.param(lambda c: c.async_increment("lit7520", 1.0), id="async_increment"), pytest.param( From d13e8dcae23326fdf3a331e5078c75d8d88c0ff7 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 09:19:43 +0000 Subject: [PATCH 087/112] fix(utils): stop wrapper_async submitting the sync success handler twice _client_async_logging_helper re-submitted logging_obj.success_handler to the executor after _dispatch_success_logging had already done so, running the same success pipeline twice per async request and racing on shared logging state. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/utils.py | 9 ------ tests/test_litellm/test_utils.py | 52 ++++++++++++++++++++++++++++++-- 2 files changed, 50 insertions(+), 11 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index 04139a124b6..4d46083e96d 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1260,15 +1260,6 @@ async def _client_async_logging_helper( async_coroutine=logging_obj.async_success_handler(result=result, start_time=start_time, end_time=end_time) ) - ################################################ - # Sync Logging Worker - ################################################ - logging_obj.handle_sync_success_callbacks_for_async_calls( - result=result, - start_time=start_time, - end_time=end_time, - ) - def _get_wrapper_num_retries(kwargs: dict[str, Any], exception: Exception) -> tuple[int | None, dict[str, Any]]: """ diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 2ace005a8e2..23bfd46bcae 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -7,8 +7,8 @@ import os import queue import threading from datetime import datetime, timedelta, timezone -from collections.abc import Iterator -from concurrent.futures import ThreadPoolExecutor +from collections.abc import Callable, Iterator +from concurrent.futures import Future, ThreadPoolExecutor from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -29,6 +29,7 @@ from litellm._logging import ( verbose_logger, ) from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.thread_pool_executor import executor as logging_executor from litellm.proxy.utils import is_valid_api_key from litellm.types.utils import ( CallTypes, @@ -6444,6 +6445,53 @@ async def test_acompletion_finishes_response_metadata_before_handing_the_respons assert snapshot["api_base"] +class _GatedSyncLoggingHookRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.seen: Final = queue.SimpleQueue[str | None]() + self.release: Final = threading.Event() + + def logging_hook( + self, kwargs: dict[str, object], result: object, call_type: str + ) -> tuple[dict[str, object], object]: + self.seen.put(result.id if isinstance(result, litellm.ModelResponse) else None) + self.release.wait(timeout=5) + return kwargs, result + + +@pytest.mark.asyncio +async def test_acompletion_runs_a_custom_logger_sync_logging_hook_exactly_once(monkeypatch: pytest.MonkeyPatch) -> None: + def legacy_sync_callback( + kwargs: dict[str, object], response: litellm.ModelResponse, start_time: datetime, end_time: datetime + ) -> None: + pass + + recorder: Final = _GatedSyncLoggingHookRecorder() + monkeypatch.setattr(litellm, "success_callback", [legacy_sync_callback, recorder]) + logging_futures: Final = queue.SimpleQueue[Future[object]]() + real_submit: Final = logging_executor.submit + + def submit_and_track(fn: Callable[..., object], *args: object, **kwargs: object) -> Future[object]: + future: Final = real_submit(fn, *args, **kwargs) + logging_futures.put(future) + return future + + with patch( # test-quality-ok: wraps the real submit only to collect the futures to join, the pool still runs + "litellm.litellm_core_utils.litellm_logging.executor.submit", side_effect=submit_and_track + ): + response: Final = await litellm.acompletion( + model="gpt-5.5", + messages=[{"role": "user", "content": "hi"}], + mock_response="Hello there!", + num_retries=0, + ) + await asyncio.sleep(0) + recorder.release.set() + for _ in range(logging_futures.qsize()): + logging_futures.get_nowait().result(timeout=5) + assert [recorder.seen.get_nowait() for _ in range(recorder.seen.qsize())] == [response.id] + + def test_completion_finishes_response_metadata_before_handing_the_response_to_the_logging_thread(): with _recording_hidden_params_at_submit("litellm.utils.executor.submit") as seen: litellm.completion( From 39f6ac4788a707012f67718f402e94e45d53cf16 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 19:47:51 +0000 Subject: [PATCH 088/112] perf(proxy): serialize /model/info listing once with orjson Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 15 +++- tests/proxy_unit_tests/test_proxy_server.py | 6 +- .../test_team_model_name_translation.py | 24 ++--- .../proxy/test_model_info_default_limits.py | 11 ++- .../proxy/test_model_list_healthy_only.py | 9 +- tests/test_litellm/proxy/test_proxy_server.py | 89 ++++++++++++++++++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 +- 7 files changed, 129 insertions(+), 29 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 43d2a149758..e7b8dd01dae 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -15036,6 +15036,13 @@ def _get_proxy_model_info(model: dict) -> dict: return _translate_model_name_for_response(model) +def _model_info_json_response(data: Sequence[Mapping[str, object]] | Mapping[str, object]) -> Response: + return Response( + content=orjson.dumps({"data": data}, default=jsonable_encoder, option=orjson.OPT_NON_STR_KEYS), + media_type="application/json", + ) + + @router.get( "/model/info", tags=["model management"], @@ -15083,7 +15090,7 @@ async def model_info_v1( `model_info.direct_access` when the proxy database is connected. Returns: - Returns a dictionary containing information about each model. + A JSON response whose `data` list holds one entry per model. Example Response: ```json @@ -15131,7 +15138,7 @@ async def model_info_v1( deployment_dict=_deployment_info_dict, excluded_keys={"litellm_credential_name"}, ) - return {"data": _deployment_info_dict} + return _model_info_json_response(_deployment_info_dict) if llm_model_list is None: raise HTTPException( @@ -15182,7 +15189,7 @@ async def model_info_v1( llm_router=llm_router, user_api_key_dict=user_api_key_dict, ) - return {"data": single_model_list} + return _model_info_json_response(single_model_list) # Return router deployments (same source as /v2/model/info), not wildcard- # expanded model names from get_complete_model_list(). Team-scoped rows @@ -15250,7 +15257,7 @@ async def model_info_v1( visible_models: Final = [model for model in all_models if model.get("model_name") not in hidden_names] verbose_proxy_logger.debug("all_models: %s", visible_models) - return {"data": visible_models} + return _model_info_json_response(visible_models) @router.get( diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index d06eb0426c9..9c8dd90dd2b 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -2139,7 +2139,7 @@ async def test_model_info_alias_without_prisma(hidden): user_api_key_dict=UserAPIKeyAuth(models=[]), ) - models = resp["data"] + models = json.loads(resp.body)["data"] alias_found = any( m["model_name"] == model_alias @@ -2203,7 +2203,7 @@ async def test_proxy_model_group_alias_checks(prisma_client, hidden): # noqa: F resp = await model_info_v1( user_api_key_dict=UserAPIKeyAuth(models=[]), ) - models = resp["data"] + models = json.loads(resp.body)["data"] is_model_alias_in_list = False for item in models: if model_alias == item["model_name"]: @@ -2280,7 +2280,7 @@ async def test_proxy_model_group_info_rerank(prisma_client): # noqa: F811 # py resp = await model_info_v1( user_api_key_dict=UserAPIKeyAuth(models=[]), ) - models = resp["data"] + models = json.loads(resp.body)["data"] assert models[0]["model_info"]["mode"] == "rerank" resp = await model_group_info( user_api_key_dict=UserAPIKeyAuth(models=[]), diff --git a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py index bc346874e0d..baa032f75e6 100644 --- a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py +++ b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py @@ -9,6 +9,7 @@ rows instead of the internal routing key `model_name_{team_id}_{uuid}`. from __future__ import annotations +import json from unittest.mock import AsyncMock, MagicMock import pytest @@ -238,7 +239,7 @@ async def test_model_info_v1_list_path_translates_team_model_name(monkeypatch): ) resp = await ps.model_info_v1(user_api_key_dict=admin, litellm_model_id=None) - names = [m["model_name"] for m in resp["data"]] + names = [m["model_name"] for m in json.loads(resp.body)["data"]] assert "team-claude-sonnet" in names assert "model_name_team-abc-123_4a6b8" not in names @@ -271,7 +272,7 @@ async def test_model_info_v1_unrestricted_key_returns_all_deployments(monkeypatc ) resp = await ps.model_info_v1(user_api_key_dict=caller, litellm_model_id=None) - assert [m["model_name"] for m in resp["data"]] == ["gpt-4"] + assert [m["model_name"] for m in json.loads(resp.body)["data"]] == ["gpt-4"] @pytest.mark.asyncio @@ -303,7 +304,7 @@ async def test_model_info_v1_restricted_key_filters_deployments(monkeypatch): ) resp = await ps.model_info_v1(user_api_key_dict=caller, litellm_model_id=None) - assert [m["model_name"] for m in resp["data"]] == ["gpt-4"] + assert [m["model_name"] for m in json.loads(resp.body)["data"]] == ["gpt-4"] def _other_team_row() -> dict: @@ -367,10 +368,11 @@ async def test_model_info_v1_unrestricted_key_hides_other_team_byok(monkeypatch) ) resp = await ps.model_info_v1(user_api_key_dict=caller, litellm_model_id=None) - returned_ids = {m["model_info"]["id"] for m in resp["data"]} + data = json.loads(resp.body)["data"] + returned_ids = {m["model_info"]["id"] for m in data} assert returned_ids == {"global-id-1", "byok-id-1"} assert "byok-id-other" not in returned_ids - names = [m["model_name"] for m in resp["data"]] + names = [m["model_name"] for m in data] assert "team-claude-sonnet" in names assert "gpt-4" in names @@ -412,7 +414,7 @@ async def test_model_info_v1_service_key_hides_all_team_byok(monkeypatch): ) resp = await ps.model_info_v1(user_api_key_dict=caller, litellm_model_id=None) - assert [m["model_info"]["id"] for m in resp["data"]] == ["global-id-1"] + assert [m["model_info"]["id"] for m in json.loads(resp.body)["data"]] == ["global-id-1"] @pytest.mark.asyncio @@ -466,7 +468,7 @@ async def test_model_info_v1_team_key_sees_own_byok_regardless_of_user_lookup( ) resp = await ps.model_info_v1(user_api_key_dict=caller, litellm_model_id=None) - assert [m["model_info"]["id"] for m in resp["data"]] == ["byok-id-1", "global-id-1"] + assert [m["model_info"]["id"] for m in json.loads(resp.body)["data"]] == ["byok-id-1", "global-id-1"] @pytest.mark.asyncio @@ -509,7 +511,7 @@ async def test_model_info_v1_user_team_membership_grants_byok(monkeypatch): ) resp = await ps.model_info_v1(user_api_key_dict=caller, litellm_model_id=None) - assert [m["model_info"]["id"] for m in resp["data"]] == [ + assert [m["model_info"]["id"] for m in json.loads(resp.body)["data"]] == [ "byok-id-other", "global-id-1", ] @@ -557,7 +559,7 @@ async def test_model_info_v1_populates_access_via_team_ids(monkeypatch): ) resp = await ps.model_info_v1(user_api_key_dict=admin, litellm_model_id=None) - by_id = {m["model_info"]["id"]: m for m in resp["data"]} + by_id = {m["model_info"]["id"]: m for m in json.loads(resp.body)["data"]} assert by_id["byok-id-1"]["model_info"]["access_via_team_ids"] == [team_id] assert by_id["byok-id-1"]["model_info"]["direct_access"] is False assert by_id["global-id-1"]["model_info"]["direct_access"] is True @@ -816,7 +818,7 @@ async def test_model_info_v1_litellm_model_id_include_team_models_filters_inacce include_team_models=True, ) - assert resp["data"] == [] + assert json.loads(resp.body)["data"] == [] @pytest.mark.asyncio @@ -852,7 +854,7 @@ async def test_model_info_v1_litellm_model_id_team_id_applies_team_filter(monkey teamId="other-team", ) - assert resp["data"] == [] + assert json.loads(resp.body)["data"] == [] team_filter.assert_awaited_once() assert team_filter.await_args.kwargs["team_id"] == "other-team" assert team_filter.await_args.kwargs["all_models"] == [team_row] diff --git a/tests/test_litellm/proxy/test_model_info_default_limits.py b/tests/test_litellm/proxy/test_model_info_default_limits.py index 8111a7af006..4426252cb6c 100644 --- a/tests/test_litellm/proxy/test_model_info_default_limits.py +++ b/tests/test_litellm/proxy/test_model_info_default_limits.py @@ -3,6 +3,7 @@ Tests verifying that default_api_key_tpm_limit and default_api_key_rpm_limit set litellm_params are returned by the /model/info endpoint. """ +import json from typing import Optional from unittest.mock import MagicMock, patch @@ -128,8 +129,9 @@ class TestModelInfoEndpointWithRouter: litellm_model_id="some-model-id", ) - assert len(response["data"]) == 1 - litellm_params = response["data"][0]["litellm_params"] + data = json.loads(response.body)["data"] + assert len(data) == 1 + litellm_params = data[0]["litellm_params"] assert litellm_params.get("default_api_key_tpm_limit") == 100 assert litellm_params.get("default_api_key_rpm_limit") == 200 @@ -171,7 +173,8 @@ class TestModelInfoEndpointWithRouter: litellm_model_id=None, ) - assert len(response["data"]) >= 1 - litellm_params = response["data"][0]["litellm_params"] + data = json.loads(response.body)["data"] + assert len(data) >= 1 + litellm_params = data[0]["litellm_params"] assert litellm_params.get("default_api_key_tpm_limit") == 100 assert litellm_params.get("default_api_key_rpm_limit") == 200 diff --git a/tests/test_litellm/proxy/test_model_list_healthy_only.py b/tests/test_litellm/proxy/test_model_list_healthy_only.py index 03eaa2e79c9..718c7e41da8 100644 --- a/tests/test_litellm/proxy/test_model_list_healthy_only.py +++ b/tests/test_litellm/proxy/test_model_list_healthy_only.py @@ -6,6 +6,7 @@ per-request `healthy_only` query parameter and the proxy-wide (`model_info_v1`). """ +import json from unittest.mock import AsyncMock, MagicMock import pytest @@ -275,7 +276,7 @@ async def test_model_info_v1_healthy_only_hides_unhealthy_deployments( litellm_model_id=None, healthy_only=True, ) - assert [m["model_name"] for m in response["data"]] == ["gpt-4"] + assert [m["model_name"] for m in json.loads(response.body)["data"]] == ["gpt-4"] @pytest.mark.asyncio @@ -286,7 +287,7 @@ async def test_model_info_v1_general_setting_hides_unhealthy_deployments(patched user_api_key_dict=_admin_key(), litellm_model_id=None, ) - assert [m["model_name"] for m in response["data"]] == ["gpt-4"] + assert [m["model_name"] for m in json.loads(response.body)["data"]] == ["gpt-4"] @pytest.mark.asyncio @@ -297,7 +298,7 @@ async def test_model_info_v1_default_keeps_unhealthy_deployments( user_api_key_dict=_admin_key(), litellm_model_id=None, ) - assert [m["model_name"] for m in response["data"]] == ["gpt-4", "claude-sonnet"] + assert [m["model_name"] for m in json.loads(response.body)["data"]] == ["gpt-4", "claude-sonnet"] patched_model_info_v1.async_get_fully_unhealthy_model_names.assert_not_awaited() @@ -318,4 +319,4 @@ async def test_model_info_v1_litellm_model_id_lookup_ignores_health_filter(patch user_api_key_dict=_admin_key(), litellm_model_id="unhealthy-id", ) - assert [m["model_name"] for m in response["data"]] == ["claude-sonnet"] + assert [m["model_name"] for m in json.loads(response.body)["data"]] == ["claude-sonnet"] diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 03a24ec5e98..a5e6cf285ad 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -15,10 +15,12 @@ from unittest import mock from unittest.mock import AsyncMock, MagicMock, create_autospec, mock_open, patch import click +import fastapi.routing import httpx import pytest import yaml from fastapi import FastAPI +from fastapi.encoders import jsonable_encoder from fastapi.staticfiles import StaticFiles from fastapi.testclient import TestClient @@ -5247,6 +5249,8 @@ async def test_model_info_v1_oci_secrets_not_leaked(): result = await model_info_v1(user_api_key_dict=mock_user_api_key_dict, litellm_model_id=None) # Verify the result structure + result_str = result.body.decode() + result = json.loads(result_str) assert "data" in result assert len(result["data"]) == 1 @@ -5269,13 +5273,96 @@ async def test_model_info_v1_oci_secrets_not_leaked(): assert litellm_params["model"].startswith("oci/"), "model should retain its full value" # Verify that actual secret values are not present in the response - result_str = str(result) assert "ocid1.api_key.oc1..aaaaaaaa7kbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbk" not in result_str assert "aa:bb:cc:dd:ee:ff:11:22:33:44:55:66:77:88:99:00" not in result_str assert "ocid1.tenancy.oc1..aaaaaaaa7kbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbk" not in result_str assert "/path/to/oci_api_key.pem" not in result_str +def test_model_info_v1_list_skips_fastapi_jsonable_encoder(monkeypatch): + """ + /model/info serializes its multi-megabyte listing itself with orjson. FastAPI must not + re-walk the payload through `jsonable_encoder`, while values orjson cannot encode natively + still come out as JSON. + """ + created_at = datetime(2026, 1, 2, 3, 4, 5, tzinfo=timezone.utc) + model_data = { + "model_name": "gpt-4o", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-secret-value"}, + "model_info": { + "id": "db-row-1", + "db_model": True, + "created_at": created_at, + "supported_regions": frozenset({"eu"}), + }, + } + mock_router = MagicMock() + mock_router.model_list = [model_data] + mock_router.get_model_list_from_model_alias.return_value = [] + mock_router.get_model_names.return_value = ["gpt-4o"] + mock_router.get_model_access_groups.return_value = {} + mock_router.get_deployment.return_value = None + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_router) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_model_list", [model_data]) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"infer_model_from_keys": False}) + monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None) + + encoder_spy = MagicMock(wraps=jsonable_encoder) + monkeypatch.setattr(fastapi.routing, "jsonable_encoder", encoder_spy) + + original_overrides = app.dependency_overrides.copy() + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1234", models=[], team_models=[] + ) + client = TestClient(app) + try: + response = client.get("/model/info") + finally: + app.dependency_overrides = original_overrides + + assert response.status_code == 200 + assert response.headers["content-type"] == "application/json" + rows = response.json()["data"] + assert [row["model_name"] for row in rows] == ["gpt-4o"] + assert rows[0]["model_info"]["created_at"] == created_at.isoformat() + assert rows[0]["model_info"]["supported_regions"] == ["eu"] + assert "sk-secret-value" not in response.text + assert encoder_spy.call_count == 0 + + +def test_model_info_v1_cli_model_returns_single_deployment_as_json(monkeypatch): + """ + A proxy started with `litellm --model ` answers /model/info with one deployment + object under `data`, serialized the same way as the listing. + """ + monkeypatch.setattr("litellm.proxy.proxy_server.user_model", "gpt-4o") + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_model_list", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + + encoder_spy = MagicMock(wraps=jsonable_encoder) + monkeypatch.setattr(fastapi.routing, "jsonable_encoder", encoder_spy) + + original_overrides = app.dependency_overrides.copy() + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1234", models=[], team_models=[] + ) + client = TestClient(app) + try: + response = client.get("/model/info") + finally: + app.dependency_overrides = original_overrides + + assert response.status_code == 200 + assert response.headers["content-type"] == "application/json" + deployment = response.json()["data"] + assert deployment["model_name"] == "*" + assert deployment["litellm_params"]["model"] == "gpt-4o" + assert encoder_spy.call_count == 0 + + def test_add_callback_from_db_to_in_memory_litellm_callbacks(): """ Test that _add_callback_from_db_to_in_memory_litellm_callbacks correctly adds callbacks diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 8c4ed5d7ff5..0b0e3e18215 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -8938,7 +8938,7 @@ export interface paths { * `model_info.direct_access` when the proxy database is connected. * * Returns: - * Returns a dictionary containing information about each model. + * A JSON response whose `data` list holds one entry per model. * * Example Response: * ```json @@ -19084,7 +19084,7 @@ export interface paths { * `model_info.direct_access` when the proxy database is connected. * * Returns: - * Returns a dictionary containing information about each model. + * A JSON response whose `data` list holds one entry per model. * * Example Response: * ```json From caaf368652f87f73277faa6353cd2e0b4261de6e Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 19:48:45 +0000 Subject: [PATCH 089/112] fix(proxy): move credentials hint helper into discovery module to break CodeQL import cycle Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/auth_utils.py | 8 -------- .../discovery_endpoints/ui_discovery_endpoints.py | 11 ++++++++++- litellm/proxy/management_endpoints/ui_sso.py | 2 +- litellm/proxy/proxy_server.py | 2 +- 4 files changed, 12 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 128b789ee75..be65c3b39ec 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -552,14 +552,6 @@ def _coerce_metadata_to_dict(value: Any) -> dict[str, Any] | None: return None -def should_hide_default_credentials_hint(general_settings: Mapping[str, object]) -> bool: - return ( - os.getenv("LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT", "false").lower() == "true" - or general_settings.get("hide_default_credentials_hint", False) is True - or bool(os.getenv("UI_PASSWORD")) - ) - - async def pre_db_read_auth_checks( request: Request, request_data: dict, diff --git a/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py b/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py index 07fd04a74fb..e0efe9dea2c 100644 --- a/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py +++ b/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py @@ -1,5 +1,6 @@ #### Analytics Endpoints ##### import os +from collections.abc import Mapping from typing import Final from fastapi import APIRouter @@ -11,10 +12,18 @@ from litellm.types.proxy.discovery_endpoints.ui_discovery_endpoints import ( router: Final = APIRouter() +def should_hide_default_credentials_hint(general_settings: Mapping[str, object]) -> bool: + return ( + os.getenv("LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT", "false").lower() == "true" + or general_settings.get("hide_default_credentials_hint", False) is True + or bool(os.getenv("UI_PASSWORD")) + ) + + @router.get("/.well-known/litellm-ui-config", response_model=UiDiscoveryEndpoints) @router.get("/litellm/.well-known/litellm-ui-config", response_model=UiDiscoveryEndpoints) # if mounted at root path async def get_ui_config(): - from litellm.proxy.auth.auth_utils import has_user_setup_sso, should_hide_default_credentials_hint + from litellm.proxy.auth.auth_utils import has_user_setup_sso from litellm.proxy.proxy_server import general_settings from litellm.proxy.utils import get_proxy_base_url, get_server_root_path diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index ab916e8df1b..6a775998f4f 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -92,7 +92,6 @@ from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken, get_user_obje from litellm.proxy.auth.auth_utils import ( _get_request_ip_address, has_user_setup_sso, - should_hide_default_credentials_hint, ) from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.ip_address_utils import IPAddressUtils @@ -107,6 +106,7 @@ from litellm.proxy.common_utils.html_forms.jwt_display_template import ( ) from litellm.proxy.common_utils.html_forms.ui_login import build_ui_login_form from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.discovery_endpoints.ui_discovery_endpoints import should_hide_default_credentials_hint from litellm.proxy.management_endpoints.internal_user_endpoints import new_user from litellm.proxy.management_endpoints.sso import CustomMicrosoftSSO from litellm.proxy.management_endpoints.sso.id_jag_assertion_capture import ( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 5b3d6365df4..230513d7eb4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -318,7 +318,6 @@ from litellm.proxy.auth.auth_utils import ( check_response_size_is_safe, is_request_body_safe, log_once_if_budget_reservation_disabled, - should_hide_default_credentials_hint, warn_once_if_custom_auth_skips_common_checks, ) from litellm.proxy.auth.fallback_model_access import router_fallback_access_check @@ -451,6 +450,7 @@ from litellm.proxy.discovery_endpoints import ( agent_skills_discovery_router, ui_discovery_endpoints_router, ) +from litellm.proxy.discovery_endpoints.ui_discovery_endpoints import should_hide_default_credentials_hint from litellm.proxy.fine_tuning_endpoints.endpoints import router as fine_tuning_router from litellm.proxy.fine_tuning_endpoints.endpoints import set_fine_tuning_config from litellm.proxy.google_endpoints.endpoints import router as google_router From 81ffc3125f989eed375e1d5a708f2a8a60cbfda1 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 19:50:01 +0000 Subject: [PATCH 090/112] fix(proxy): release realtime max_parallel_requests slot when the task is cancelled during pre-call Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 115 +++++++++--------- tests/test_litellm/proxy/test_proxy_server.py | 27 ++-- 2 files changed, 77 insertions(+), 65 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ed3d43800f7..8f046f0892b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11912,7 +11912,6 @@ async def _reject_realtime_session( await websocket.close(code=code, reason=reason) finally: await _release_realtime_budget_reservation(user_api_key_dict) - await proxy_logging_obj._arelease_max_parallel_requests_on_disconnect(user_api_key_dict) # pyright: ignore[reportPrivateUsage] # same release idiom the HTTP disconnect path uses @app.websocket("/openai/v1/realtime") @@ -11992,65 +11991,67 @@ async def realtime_websocket_endpoint( # Errors here (e.g. guardrail block) are sent back to the client as an # error event before closing, so the caller knows what happened. try: - ( - data, - litellm_logging_obj, - ) = await base_llm_response_processor.common_processing_pre_call_logic( - request=request, - general_settings=general_settings, - user_api_key_dict=user_api_key_dict, - version=version, - proxy_logging_obj=proxy_logging_obj, - proxy_config=proxy_config, - user_model=user_model, - user_temperature=user_temperature, - user_request_timeout=user_request_timeout, - user_max_tokens=user_max_tokens, - user_api_base=user_api_base, - model=route_model, - route_type="_arealtime", - ) - except Exception as e: - verbose_proxy_logger.exception("Realtime pre-call error") - await _reject_realtime_session( - websocket, user_api_key_dict, code=1011, reason="Pre-call error", error_message=str(e) - ) - return - - # Phase 2: route to upstream LLM. - try: - data["user_api_key_dict"] = user_api_key_dict - llm_call: Final = await route_request( - data=data, - route_type="_arealtime", - llm_router=llm_router, - user_model=user_model, - ) - await llm_call - except websockets.exceptions.InvalidStatusCode as e: - verbose_proxy_logger.exception("Invalid status code") - await websocket.close(code=e.status_code, reason="Invalid status code") - except Exception as e: - verbose_proxy_logger.exception("Internal server error") - redacted_error: Final = _redact_string(str(e)) try: - await websocket.send_text(realtime_error_event(redacted_error, error_type="server_error")) - except Exception: # noqa: BLE001 # best-effort notice: a dead client socket must not skip the close below - verbose_proxy_logger.debug("Could not send realtime error event to client; closing anyway") - try: - await websocket.close( - code=1011, - reason=websocket_close_reason(redacted_error, fallback="Internal server error"), + ( + data, + litellm_logging_obj, + ) = await base_llm_response_processor.common_processing_pre_call_logic( + request=request, + general_settings=general_settings, + user_api_key_dict=user_api_key_dict, + version=version, + proxy_logging_obj=proxy_logging_obj, + proxy_config=proxy_config, + user_model=user_model, + user_temperature=user_temperature, + user_request_timeout=user_request_timeout, + user_max_tokens=user_max_tokens, + user_api_base=user_api_base, + model=route_model, + route_type="_arealtime", ) - except Exception: # noqa: BLE001 # the lower layer may have closed the socket already; closing twice is not an error - verbose_proxy_logger.debug("Could not close realtime client websocket; it is already gone") - finally: - from litellm.litellm_core_utils.realtime_streaming import ( - REALTIME_SESSION_SUCCESS_LOGGED_KEY, - ) + except Exception as e: + verbose_proxy_logger.exception("Realtime pre-call error") + await _reject_realtime_session( + websocket, user_api_key_dict, code=1011, reason="Pre-call error", error_message=str(e) + ) + return - if not litellm_logging_obj.model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY): - await _release_realtime_budget_reservation(user_api_key_dict) + # Phase 2: route to upstream LLM. + try: + data["user_api_key_dict"] = user_api_key_dict + llm_call: Final = await route_request( + data=data, + route_type="_arealtime", + llm_router=llm_router, + user_model=user_model, + ) + await llm_call + except websockets.exceptions.InvalidStatusCode as e: + verbose_proxy_logger.exception("Invalid status code") + await websocket.close(code=e.status_code, reason="Invalid status code") + except Exception as e: + verbose_proxy_logger.exception("Internal server error") + redacted_error: Final = _redact_string(str(e)) + try: + await websocket.send_text(realtime_error_event(redacted_error, error_type="server_error")) + except Exception: # noqa: BLE001 # best-effort notice: a dead client socket must not skip the close below + verbose_proxy_logger.debug("Could not send realtime error event to client; closing anyway") + try: + await websocket.close( + code=1011, + reason=websocket_close_reason(redacted_error, fallback="Internal server error"), + ) + except Exception: # noqa: BLE001 # the lower layer may have closed the socket already; closing twice is not an error + verbose_proxy_logger.debug("Could not close realtime client websocket; it is already gone") + finally: + from litellm.litellm_core_utils.realtime_streaming import ( + REALTIME_SESSION_SUCCESS_LOGGED_KEY, + ) + + if not litellm_logging_obj.model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY): + await _release_realtime_budget_reservation(user_api_key_dict) + finally: await proxy_logging_obj._arelease_max_parallel_requests_on_disconnect(user_api_key_dict) # pyright: ignore[reportPrivateUsage] # same release idiom the HTTP disconnect path uses diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 7b6a838aa3e..11c2d301365 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -9919,8 +9919,9 @@ async def _lit6973_drive_realtime_session( phase_one_exit picks a rejection before the relay: "model_access" makes the key/model check raise ProxyException, "pre_call" makes pre-call processing - (rate limits, guardrails) raise. Neither reaches route_request, so no success - log can own the reservation and the endpoint has to release it on that exit. + (rate limits, guardrails) raise, "pre_call_cancelled" cancels the task inside + pre-call processing. None reaches route_request, so no success log can own the + reservation and the endpoint has to release it on that exit. route_request resolves normally in both cases: the relay owns the session once route_request returns. A successful session enqueues its success cost @@ -9950,7 +9951,13 @@ async def _lit6973_drive_realtime_session( if phase_one_exit == "model_access" else None ) - pre_call_error: Final = Exception("Rate limit exceeded") if phase_one_exit == "pre_call" else None + pre_call_error: Final = ( + asyncio.CancelledError() + if phase_one_exit == "pre_call_cancelled" + else Exception("Rate limit exceeded") + if phase_one_exit == "pre_call" + else None + ) pre_call: Final = AsyncMock( side_effect=pre_call_error, return_value=({"model": "vertex_ai/gemini-live-2.5-flash"}, logging_obj) ) @@ -10069,15 +10076,16 @@ async def test_successful_realtime_session_leaves_the_reservation_for_the_cost_c @pytest.mark.asyncio -@pytest.mark.parametrize("phase_one_exit", [None, "pre_call"]) +@pytest.mark.parametrize("phase_one_exit", [None, "pre_call", "pre_call_cancelled"]) async def test_realtime_session_ending_without_llm_callbacks_releases_the_max_parallel_slot( phase_one_exit: str | None, ): """The rate limiter acquires the key's max_parallel_requests slot in pre-call and only frees it from the LLM success/failure callbacks. A realtime session that ends - without either callback (Bedrock closes without usage events, or a later pre-call - hook rejects the session) has to be released by the route itself, or the slot stays - occupied until its TTL and the key's next session is refused with a 429.""" + without either callback (Bedrock closes without usage events, a later pre-call hook + rejects the session, or the task is cancelled while still in pre-call) has to be + released by the route itself, or the slot stays occupied until its TTL and the key's + next session is refused with a 429.""" from litellm.caching.caching import DualCache from litellm.proxy import proxy_server as ps from litellm.proxy.hooks.parallel_request_limiter_v3 import ( @@ -10097,7 +10105,10 @@ async def test_realtime_session_ending_without_llm_callbacks_releases_the_max_pa stash_token: Final = _request_stash.set(stash) try: hooks: Final = patch.dict(ps.proxy_logging_obj.proxy_hook_mapping, {"parallel_request_limiter": limiter}) # test-quality-ok: registers a real limiter on the module-global hook map the route reads; assertion observes its counter - with hooks: + expected_exit: Final = ( + pytest.raises(asyncio.CancelledError) if phase_one_exit == "pre_call_cancelled" else contextlib.nullcontext() + ) + with hooks, expected_exit: await _lit6973_drive_realtime_session( reservation, backend_logged_success=False, phase_one_exit=phase_one_exit ) From 845596063d257482652bbc56a89d437435d73972 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 09:25:47 +0000 Subject: [PATCH 091/112] fix(prometheus): label pre-call rate limit failures with the resolved api_provider Pre-call limiters reject before a deployment is attached to request_data, so the failure hook could not resolve api_provider for router aliases and emitted api_provider="None" on litellm_proxy_failed_requests_metric_total and litellm_proxy_total_requests_metric_total. Fall back to the provider the limiter already resolved onto RateLimitError.llm_provider, keeping request data as the first source and ignoring the proxy placeholder. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/prometheus.py | 14 +++- .../integrations/test_prometheus_labels.py | 71 +++++++++++++++++++ 2 files changed, 84 insertions(+), 1 deletion(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 2528f07f92c..2ba5d471e75 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -41,6 +41,7 @@ from litellm.proxy._types import ( LiteLLM_UserTable, UserAPIKeyAuth, ) +from litellm.proxy.hooks.rate_limiter_utils import PROXY_LLM_PROVIDER_FALLBACK from litellm.repositories.base_repository import BaseRepository from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.team_repository import TeamRepository @@ -2581,6 +2582,15 @@ class PrometheusLogger(CustomLogger): ) return None + @staticmethod + def _extract_api_provider_from_exception(exception: Exception) -> str | None: + if not isinstance(exception, litellm.exceptions.RateLimitError): + return None + llm_provider: Final = exception.llm_provider + if not llm_provider or llm_provider == PROXY_LLM_PROVIDER_FALLBACK: + return None + return llm_provider + async def async_post_call_failure_hook( self, request_data: dict, @@ -2616,7 +2626,9 @@ class PrometheusLogger(CustomLogger): _metadata: Final = request_data.get("metadata", {}) or {} model_id: Final = _metadata.get("model_info", {}).get("id") or request_data.get("model_info", {}).get("id") rate_limit_category, rate_limit_type = self._extract_rate_limit_labels(original_exception) - api_provider: Final = self._extract_api_provider_from_request_data(request_data) + api_provider: Final = self._extract_api_provider_from_request_data( + request_data + ) or self._extract_api_provider_from_exception(original_exception) enum_values: Final = UserAPIKeyLabelValues( end_user=user_api_key_dict.end_user_id, user=user_api_key_dict.user_id, diff --git a/tests/test_litellm/integrations/test_prometheus_labels.py b/tests/test_litellm/integrations/test_prometheus_labels.py index 859cdd30c11..200f8e65add 100644 --- a/tests/test_litellm/integrations/test_prometheus_labels.py +++ b/tests/test_litellm/integrations/test_prometheus_labels.py @@ -716,6 +716,77 @@ async def test_failure_hook_emits_api_provider_value_on_failed_requests_metric() _clear_prometheus_registry() +async def _failed_requests_api_provider_labels( + request_data: dict[str, object], + original_exception: Exception, +) -> list[str]: + from litellm.integrations.prometheus import PrometheusLogger + from litellm.proxy._types import UserAPIKeyAuth + + _clear_prometheus_registry() + try: + await PrometheusLogger().async_post_call_failure_hook( + request_data=request_data, + original_exception=original_exception, + user_api_key_dict=UserAPIKeyAuth(token="tok"), + ) + return [ + s.labels.get("api_provider") + for s in _collected_samples("litellm_proxy_failed_requests_metric_total") + ] + finally: + _clear_prometheus_registry() + + +@pytest.mark.asyncio +async def test_failure_hook_emits_api_provider_from_pre_call_rate_limit_error_for_router_alias(): + """ + Pre-call limiters reject before a deployment lands on request_data and a + router alias cannot be inferred from its name, so the provider the limiter + resolved onto the exception is the only source for the label. + """ + from litellm.exceptions import RateLimitType + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + + err = ProxyRateLimitError( + detail={"error": "rpm exceeded"}, + rate_limit_type=RateLimitType.REQUESTS, + model="openai/gpt-5.4-mini", + llm_provider="openai", + ) + + assert await _failed_requests_api_provider_labels( + {"model": "team-chat-model", "metadata": {}}, err + ) == ["openai"] + + +@pytest.mark.asyncio +async def test_failure_hook_leaves_api_provider_unset_when_rate_limiter_could_not_resolve_provider(): + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + + err = ProxyRateLimitError(detail={"error": "rpm exceeded"}, model="unknown-alias") + + assert await _failed_requests_api_provider_labels( + {"model": "unknown-alias", "metadata": {}}, err + ) == ["None"] + + +@pytest.mark.asyncio +async def test_failure_hook_prefers_request_data_provider_over_exception_provider(): + from litellm.exceptions import RateLimitError + + err = RateLimitError(message="upstream 429", llm_provider="openai", model="gpt-4o") + + assert await _failed_requests_api_provider_labels( + { + "model": "gpt-4o", + "metadata": {}, + "litellm_params": {"custom_llm_provider": "azure"}, + }, + err, + ) == ["azure"] + + if __name__ == "__main__": test_user_email_in_required_metrics() test_user_email_label_exists() From 10be7e01d4ac373a408696c26efcd1a30cb7d2d9 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 19:54:00 +0000 Subject: [PATCH 092/112] fix(router): import retry helpers inside their functions to break CodeQL import cycles Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 3 ++- litellm/utils.py | 4 +++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index a46ffa83b05..31ea81e35be 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -241,7 +241,6 @@ from litellm.types.router import ( ModelGroupInfo, OptionalPreCallChecks, PreRoutingStrategy, - RetryAttemptRecord, RetryPolicy, RouterCacheEnum, RouterErrors, @@ -8362,6 +8361,8 @@ class Router: """ When a retry or fallback happens, record which model group, deployment and attempt just failed and why """ + from litellm.types.router import RetryAttemptRecord + _metadata_var: Final = "litellm_metadata" if "litellm_metadata" in kwargs else "metadata" request_metadata: Final[Mapping[str, object]] = kwargs[_metadata_var] model_group: Final = kwargs.get("model") diff --git a/litellm/utils.py b/litellm/utils.py index a0e910a1616..cc1c52d1c78 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -81,7 +81,7 @@ from litellm.constants import ( PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO, TOOL_CHOICE_OBJECT_TOKEN_COUNT, ) -from litellm.litellm_core_utils.core_helpers import max_retries_per_request_hit, normalize_drop_params +from litellm.litellm_core_utils.core_helpers import normalize_drop_params from litellm.litellm_core_utils.fallback_generalizations import ( match_capability_generalizations, ) @@ -1500,6 +1500,8 @@ def post_call_processing( def client(original_function): + from litellm.litellm_core_utils.core_helpers import max_retries_per_request_hit + Rules: Final = litellm_utils.Rules rules_obj: Final = Rules() From 21cd52d508ea61cb82e969b2384f1353a61c89e5 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 19:55:06 +0000 Subject: [PATCH 093/112] fix(bedrock/realtime): declare the client websocket scope as a protocol attribute CodeQL py/ineffectual-statement flags the bare ellipsis body of the @property declaration on the RealtimeClientWebSocket protocol. A plain attribute annotation states the same structural contract without an expression statement. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/bedrock/realtime/handler.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 130dec20a4b..7ab5a14dfc2 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -84,8 +84,7 @@ def _should_log_event(openai_message: Mapping[str, object]) -> bool: class RealtimeClientWebSocket(Protocol): """The client-facing websocket surface the realtime bridge talks to.""" - @property - def scope(self) -> MutableMapping[str, object]: ... # mutable-ok: the ASGI scope is the per-connection state store + scope: MutableMapping[str, object] # mutable-ok: the ASGI scope is the per-connection state store async def receive_text(self) -> str: ... From 9e870ffd005ffb05357feb1328eaeefe586b3eb8 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 19:57:50 +0000 Subject: [PATCH 094/112] fix(proxy): resolve CodeQL findings on team list, sanitize log args and move UserNotFoundError out of the import cycle Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 7 +++++++ litellm/proxy/auth/auth_checks.py | 8 +------- litellm/proxy/management_endpoints/team_endpoints.py | 6 +++--- 3 files changed, 11 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 49e0247aad9..412b902febc 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3954,6 +3954,13 @@ class ManagementEndpointLoggingPayload(LiteLLMPydanticObjectBase): end_time: datetime | None = None +class UserNotFoundError(ValueError): + """The user row is provably absent, as opposed to merely unreadable, so a caller that reads a missing row as no user-level limits can key on it without also swallowing a database that would not answer.""" + + def __init__(self, user_id: str) -> None: + super().__init__(f"User doesn't exist in db. 'user_id'={user_id}. Create user via `/user/new` call.") + + class ProxyException(Exception): # NOTE: DO NOT MODIFY THIS # This is used to map exactly to OPENAI Exceptions diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 576585ee9a3..6f4e5cd3ca1 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -65,6 +65,7 @@ from litellm.proxy._types import ( RoleBasedPermissions, SpecialModelNames, UserAPIKeyAuth, + UserNotFoundError, ) from litellm.proxy.auth.budget_throttle import ( budget_throttle_percentage, @@ -2375,13 +2376,6 @@ async def _backfill_null_user_email( return updated_row -class UserNotFoundError(ValueError): - """The user row is provably absent, as opposed to merely unreadable, so a caller that reads a missing row as no user-level limits can key on it without also swallowing a database that would not answer.""" - - def __init__(self, user_id: str) -> None: - super().__init__(f"User doesn't exist in db. 'user_id'={user_id}. Create user via `/user/new` call.") - - @log_db_metrics async def get_user_object( user_id: str | None, diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index fdd44c16799..a627c1438dc 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -73,10 +73,10 @@ from litellm.proxy._types import ( TeamModelDeleteRequest, UpdateTeamRequest, UserAPIKeyAuth, + UserNotFoundError, ) from litellm.proxy.auth.auth_checks import ( OrganizationNotFoundError, - UserNotFoundError, _cache_team_object, allowed_route_check_inside_route, can_org_access_model, @@ -5107,9 +5107,9 @@ async def _enforce_list_team_v2_access( ) verbose_proxy_logger.debug( "list_team_v2: org admin access for user=%s, org_ids=%s, user_id_filter=%s", - caller_user_id, + _sanitize_for_log(caller_user_id), org_admin_org_ids, - None if is_own_query else user_id, + _sanitize_for_log(None if is_own_query else user_id), ) return None if is_own_query else user_id, org_admin_org_ids, own_team_ids From bea22df6c27b30d060e6f7e239e94396c9f38ad7 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 19:59:10 +0000 Subject: [PATCH 095/112] fix(caching): keep generic add_cache failures at ERROR unless the backend is Redis Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/caching/caching.py | 13 ++++++++--- tests/test_litellm/caching/test_caching.py | 26 ++++++++++++++++++++++ 2 files changed, 36 insertions(+), 3 deletions(-) diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index d6dd2a073af..be82f5def1f 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -679,7 +679,14 @@ class Cache: cache_key, cached_data, kwargs = self._add_cache_logic(result=result, **kwargs) self.cache.set_cache(cache_key, cached_data, **kwargs) except Exception as e: - log_redis_failure(verbose_logger, logging.ERROR, "LiteLLM Cache: exception in add_cache", e) + self._log_add_cache_failure(e) + + def _log_add_cache_failure(self, exc: Exception) -> None: + message: Final = "LiteLLM Cache: exception in add_cache" + if isinstance(self.cache, RedisCache): + log_redis_failure(verbose_logger, logging.ERROR, message, exc) + return + verbose_logger.error("%s: %s", message, exc) async def async_add_cache(self, result, dynamic_cache_object: BaseCache | None = None, **kwargs): """ @@ -698,7 +705,7 @@ class Cache: else: await self.cache.async_set_cache(cache_key, cached_data, **kwargs) except Exception as e: - log_redis_failure(verbose_logger, logging.ERROR, "LiteLLM Cache: exception in add_cache", e) + self._log_add_cache_failure(e) def _convert_to_cached_embedding( self, @@ -877,7 +884,7 @@ class Cache: else: await self.cache.async_set_cache_pipeline(cache_list=cache_list, **kwargs) except Exception as e: - log_redis_failure(verbose_logger, logging.ERROR, "LiteLLM Cache: exception in add_cache", e) + self._log_add_cache_failure(e) def should_use_cache(self, **kwargs): """ diff --git a/tests/test_litellm/caching/test_caching.py b/tests/test_litellm/caching/test_caching.py index 4d0ec0fb677..c7ec8abc31e 100644 --- a/tests/test_litellm/caching/test_caching.py +++ b/tests/test_litellm/caching/test_caching.py @@ -1,9 +1,12 @@ import logging import re +from unittest.mock import MagicMock import pytest +import litellm.caching.redis_cache as redis_cache_module from litellm.caching.caching import Cache +from litellm.caching.redis_cache import RedisCache, _RedisTimeoutLogThrottle from litellm.types.caching import LiteLLMCacheType, SemanticCacheScope from litellm.types.utils import Embedding, EmbeddingResponse, Usage @@ -53,6 +56,29 @@ def test_cache_key_debug_log_does_not_include_prompt_material(caplog): assert any(cache_key in message for message in created_cache_key_logs) +@pytest.mark.parametrize( + ("backend", "expected_level"), + [ + pytest.param(MagicMock(spec=RedisCache), logging.DEBUG, id="redis_backend_is_throttled"), + pytest.param(MagicMock(), logging.ERROR, id="other_backend_logs_every_timeout"), + ], +) +def test_add_cache_timeout_only_joins_redis_throttle_for_redis_backends(backend, expected_level, caplog, monkeypatch): + throttle = _RedisTimeoutLogThrottle(interval=5.0, clock=MagicMock(return_value=1_000.0)) + assert throttle.admit() == 0 + monkeypatch.setattr(redis_cache_module, "_redis_timeout_log_throttle", throttle) + + cache = Cache(type=LiteLLMCacheType.LOCAL) + backend.set_cache.side_effect = TimeoutError("lit7520 backend timed out") + cache.cache = backend + + with caplog.at_level(logging.DEBUG, logger="LiteLLM"): + cache.add_cache("result", model="gpt-4.1-mini", messages=[{"role": "user", "content": "hi"}]) + + records = [r for r in caplog.records if "lit7520 backend timed out" in r.getMessage()] + assert [r.levelno for r in records] == [expected_level] + + def _embedding_response(prompt_tokens, num_items): return EmbeddingResponse( model="amazon.titan-embed-image-v1", From a4e34d6e1b9905736acd5d13ea03d7e32f3e3f0e Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 20:05:14 +0000 Subject: [PATCH 096/112] test(caching): wrap the DualCache fixture line to the 120 character limit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/caching/test_dual_cache.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/test_litellm/caching/test_dual_cache.py b/tests/test_litellm/caching/test_dual_cache.py index 6f29be00b30..95395878c25 100644 --- a/tests/test_litellm/caching/test_dual_cache.py +++ b/tests/test_litellm/caching/test_dual_cache.py @@ -726,7 +726,10 @@ async def test_redis_timeouts_falling_back_to_memory_log_once_per_interval(caplo async def async_increment(self, key, value, **kwargs): raise RedisTimeoutError("Timeout reading from 127.0.0.1:6379") - cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=_TimingOutRedis()) # pyright: ignore[reportArgumentType] # duck-typed Redis double + cache = DualCache( + in_memory_cache=InMemoryCache(), + redis_cache=_TimingOutRedis(), # pyright: ignore[reportArgumentType] # duck-typed Redis double + ) increments = [RedisPipelineIncrementOperation(key="k", increment_value=1.0, ttl=60)] with caplog.at_level(logging.DEBUG, logger="LiteLLM"): From 4fee71b7ad991b2c1c1af0e6aa5b33f18ab57c97 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 20:10:36 +0000 Subject: [PATCH 097/112] fix(proxy): define UserNotFoundError in an import-free types module so no importer sits in the CodeQL cycle Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 7 ------- litellm/proxy/auth/auth_checks.py | 2 +- litellm/proxy/management_endpoints/team_endpoints.py | 2 +- litellm/types/proxy/auth/auth_checks.py | 8 ++++++++ 4 files changed, 10 insertions(+), 9 deletions(-) create mode 100644 litellm/types/proxy/auth/auth_checks.py diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 412b902febc..49e0247aad9 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3954,13 +3954,6 @@ class ManagementEndpointLoggingPayload(LiteLLMPydanticObjectBase): end_time: datetime | None = None -class UserNotFoundError(ValueError): - """The user row is provably absent, as opposed to merely unreadable, so a caller that reads a missing row as no user-level limits can key on it without also swallowing a database that would not answer.""" - - def __init__(self, user_id: str) -> None: - super().__init__(f"User doesn't exist in db. 'user_id'={user_id}. Create user via `/user/new` call.") - - class ProxyException(Exception): # NOTE: DO NOT MODIFY THIS # This is used to map exactly to OPENAI Exceptions diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 6f4e5cd3ca1..0cbb4a38aba 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -65,7 +65,6 @@ from litellm.proxy._types import ( RoleBasedPermissions, SpecialModelNames, UserAPIKeyAuth, - UserNotFoundError, ) from litellm.proxy.auth.budget_throttle import ( budget_throttle_percentage, @@ -124,6 +123,7 @@ from litellm.repositories.table_repositories import ( from litellm.repositories.team_repository import TeamRepository from litellm.repositories.user_repository import UserRepository from litellm.router import Router +from litellm.types.proxy.auth.auth_checks import UserNotFoundError from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget from litellm.utils import get_utc_datetime diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index a627c1438dc..6b16692f7ad 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -73,7 +73,6 @@ from litellm.proxy._types import ( TeamModelDeleteRequest, UpdateTeamRequest, UserAPIKeyAuth, - UserNotFoundError, ) from litellm.proxy.auth.auth_checks import ( OrganizationNotFoundError, @@ -157,6 +156,7 @@ from litellm.repositories.verification_token_repository import ( VerificationTokenRepository, ) from litellm.router import Router +from litellm.types.proxy.auth.auth_checks import UserNotFoundError from litellm.types.proxy.management_endpoints.common_daily_activity import ( SpendAnalyticsPaginatedResponse, ) diff --git a/litellm/types/proxy/auth/auth_checks.py b/litellm/types/proxy/auth/auth_checks.py new file mode 100644 index 00000000000..80c65d14113 --- /dev/null +++ b/litellm/types/proxy/auth/auth_checks.py @@ -0,0 +1,8 @@ +"""Failure values raised by `litellm/proxy/auth/auth_checks.py`. Kept free of `litellm` imports so any proxy module can import them without joining the `litellm.proxy` import cycle.""" + + +class UserNotFoundError(ValueError): + """The user row is provably absent, as opposed to merely unreadable, so a caller that reads a missing row as no user-level limits can key on it without also swallowing a database that would not answer.""" + + def __init__(self, user_id: str) -> None: + super().__init__(f"User doesn't exist in db. 'user_id'={user_id}. Create user via `/user/new` call.") From d7900df73f6e4b7a346e1a050e66a481ac68ff41 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 20:13:02 +0000 Subject: [PATCH 098/112] chore(constants): drop the restating comment above the Bedrock realtime scope keys Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 1 - 1 file changed, 1 deletion(-) diff --git a/litellm/constants.py b/litellm/constants.py index a1a737153fa..f94d92735fd 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -311,7 +311,6 @@ REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS: Final = float( # RFC 6455 caps the close frame payload at 125 bytes, 2 of which carry the status code WEBSOCKET_CLOSE_REASON_MAX_BYTES: Final = 123 -# ASGI websocket scope keys the Bedrock realtime bridge uses to carry state across router fallback attempts BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY: Final = "litellm.bedrock_realtime.pending_session_update" BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY: Final = "litellm.bedrock_realtime.session_committed" BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY: Final = "litellm.bedrock_realtime.committed_failure" From 4ac168b3fcb440a7e0af0e3b919ad6880baf2fd6 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 20:13:39 +0000 Subject: [PATCH 099/112] fix(auth): drop in-flight membership load on invalidation so it cannot repopulate the cache Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/auth_checks.py | 8 ++- .../proxy/auth/test_auth_checks.py | 50 +++++++++++++++++++ 2 files changed, 56 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index fa85400f4e0..43d840130b6 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2202,8 +2202,11 @@ async def _fetch_team_membership_from_db( where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, include={"litellm_budget_table": True}, ) + membership: Final = None if response is None else LiteLLM_TeamMembership.model_validate(response.dict()) _key: Final = team_membership_reservation_cache_key(user_id=user_id, team_id=team_id) - if response is None: + if _team_membership_inflight.get(_key) is not asyncio.current_task(): + return membership + if membership is None: await user_api_key_cache.async_set_cache( key=_key, value=NO_TEAM_MEMBERSHIP_SENTINEL, @@ -2211,7 +2214,6 @@ async def _fetch_team_membership_from_db( ) return None - membership: Final = LiteLLM_TeamMembership.model_validate(response.dict()) await user_api_key_cache.async_set_cache( key=_key, value=membership, @@ -2762,6 +2764,8 @@ async def invalidate_team_member_spend_state( publish_auth_cache_invalidation, ) + _team_membership_inflight.pop(team_membership_reservation_cache_key(user_id=user_id, team_id=team_id), None) + if new_spend is not None: from litellm.proxy.proxy_server import SPEND_DB_FLOOR_CACHE_TTL_SECONDS, spend_counter_cache diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index b843916debf..986d53b80b0 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -6493,6 +6493,56 @@ async def test_get_team_membership_coalesces_parallel_db_fetches(): mock_prisma_client.db.litellm_teammembership.find_unique.assert_awaited_once() +@pytest.mark.asyncio +async def test_get_team_membership_invalidation_mid_flight_discards_stale_load(): + from litellm.proxy._types import LiteLLM_TeamMembership + from litellm.proxy.auth.auth_checks import get_team_membership, invalidate_team_member_spend_state + from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec + from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key + + started = asyncio.Event() + release_stale = asyncio.Event() + release_fresh = asyncio.Event() + loads = iter((("budget-old", release_stale), ("budget-new", release_fresh))) + + async def _find_unique(*args, **kwargs): + budget_id, release = next(loads) + row = MagicMock() + row.dict = lambda: {"user_id": "u-inv", "team_id": "t-inv", "spend": 1.0, "budget_id": budget_id} + started.set() + await release.wait() + return row + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(side_effect=_find_unique) + cache = UserApiKeyCache() + + async def _load(): + return await get_team_membership( + user_id="u-inv", team_id="t-inv", prisma_client=mock_prisma_client, user_api_key_cache=cache + ) + + stale = asyncio.create_task(_load()) + await started.wait() + await invalidate_team_member_spend_state(user_id="u-inv", team_id="t-inv", user_api_key_cache=cache) + started.clear() + fresh = asyncio.create_task(_load()) + await asyncio.wait_for(started.wait(), timeout=2) + release_fresh.set() + fresh_result = await fresh + release_stale.set() + stale_result = await stale + + assert stale_result is not None and stale_result.budget_id == "budget-old" + assert fresh_result is not None and fresh_result.budget_id == "budget-new" + assert mock_prisma_client.db.litellm_teammembership.find_unique.await_count == 2 + cached = CacheCodec.deserialize( + await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-inv", team_id="t-inv")), + model_type=LiteLLM_TeamMembership, + ) + assert cached is not None and cached.budget_id == "budget-new" + + @pytest.mark.asyncio async def test_common_checks_calls_get_team_membership_once_per_request(): from fastapi import Request From 1b31be1a9c3b06543addec28c6342e14a53bd3db Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 20:23:23 +0000 Subject: [PATCH 100/112] fix(proxy): leave the realtime max_parallel slot to the success callback when one is enqueued Releasing the slot unconditionally from the route raced the limiter's own success handler on the logging worker: both could read the same stashed acquisition before either cleared it, and under the integer in-memory fallback that double-decrements the counter. The route now releases only on exits without a success callback (pre-call rejection, pre-call cancellation, and Phase 2 exits without the success stamp), matching the HTTP disconnect path's ownership rule. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 122 +++++++++--------- tests/test_litellm/proxy/test_proxy_server.py | 95 ++++++++++---- 2 files changed, 131 insertions(+), 86 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 8f046f0892b..f37744f4685 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11893,6 +11893,10 @@ async def _release_realtime_budget_reservation(user_api_key_dict: UserAPIKeyAuth ) +async def _release_realtime_max_parallel_slot(user_api_key_dict: UserAPIKeyAuth) -> None: + await proxy_logging_obj._arelease_max_parallel_requests_on_disconnect(user_api_key_dict) # pyright: ignore[reportPrivateUsage] # same release idiom the HTTP disconnect path uses + + async def _reject_realtime_session( websocket: WebSocket, user_api_key_dict: UserAPIKeyAuth, @@ -11912,6 +11916,7 @@ async def _reject_realtime_session( await websocket.close(code=code, reason=reason) finally: await _release_realtime_budget_reservation(user_api_key_dict) + await _release_realtime_max_parallel_slot(user_api_key_dict) @app.websocket("/openai/v1/realtime") @@ -11991,68 +11996,69 @@ async def realtime_websocket_endpoint( # Errors here (e.g. guardrail block) are sent back to the client as an # error event before closing, so the caller knows what happened. try: - try: - ( - data, - litellm_logging_obj, - ) = await base_llm_response_processor.common_processing_pre_call_logic( - request=request, - general_settings=general_settings, - user_api_key_dict=user_api_key_dict, - version=version, - proxy_logging_obj=proxy_logging_obj, - proxy_config=proxy_config, - user_model=user_model, - user_temperature=user_temperature, - user_request_timeout=user_request_timeout, - user_max_tokens=user_max_tokens, - user_api_base=user_api_base, - model=route_model, - route_type="_arealtime", - ) - except Exception as e: - verbose_proxy_logger.exception("Realtime pre-call error") - await _reject_realtime_session( - websocket, user_api_key_dict, code=1011, reason="Pre-call error", error_message=str(e) - ) - return + ( + data, + litellm_logging_obj, + ) = await base_llm_response_processor.common_processing_pre_call_logic( + request=request, + general_settings=general_settings, + user_api_key_dict=user_api_key_dict, + version=version, + proxy_logging_obj=proxy_logging_obj, + proxy_config=proxy_config, + user_model=user_model, + user_temperature=user_temperature, + user_request_timeout=user_request_timeout, + user_max_tokens=user_max_tokens, + user_api_base=user_api_base, + model=route_model, + route_type="_arealtime", + ) + except Exception as e: + verbose_proxy_logger.exception("Realtime pre-call error") + await _reject_realtime_session( + websocket, user_api_key_dict, code=1011, reason="Pre-call error", error_message=str(e) + ) + return + except BaseException: + await _release_realtime_max_parallel_slot(user_api_key_dict) + raise - # Phase 2: route to upstream LLM. + # Phase 2: route to upstream LLM. + try: + data["user_api_key_dict"] = user_api_key_dict + llm_call: Final = await route_request( + data=data, + route_type="_arealtime", + llm_router=llm_router, + user_model=user_model, + ) + await llm_call + except websockets.exceptions.InvalidStatusCode as e: + verbose_proxy_logger.exception("Invalid status code") + await websocket.close(code=e.status_code, reason="Invalid status code") + except Exception as e: + verbose_proxy_logger.exception("Internal server error") + redacted_error: Final = _redact_string(str(e)) try: - data["user_api_key_dict"] = user_api_key_dict - llm_call: Final = await route_request( - data=data, - route_type="_arealtime", - llm_router=llm_router, - user_model=user_model, + await websocket.send_text(realtime_error_event(redacted_error, error_type="server_error")) + except Exception: # noqa: BLE001 # best-effort notice: a dead client socket must not skip the close below + verbose_proxy_logger.debug("Could not send realtime error event to client; closing anyway") + try: + await websocket.close( + code=1011, + reason=websocket_close_reason(redacted_error, fallback="Internal server error"), ) - await llm_call - except websockets.exceptions.InvalidStatusCode as e: - verbose_proxy_logger.exception("Invalid status code") - await websocket.close(code=e.status_code, reason="Invalid status code") - except Exception as e: - verbose_proxy_logger.exception("Internal server error") - redacted_error: Final = _redact_string(str(e)) - try: - await websocket.send_text(realtime_error_event(redacted_error, error_type="server_error")) - except Exception: # noqa: BLE001 # best-effort notice: a dead client socket must not skip the close below - verbose_proxy_logger.debug("Could not send realtime error event to client; closing anyway") - try: - await websocket.close( - code=1011, - reason=websocket_close_reason(redacted_error, fallback="Internal server error"), - ) - except Exception: # noqa: BLE001 # the lower layer may have closed the socket already; closing twice is not an error - verbose_proxy_logger.debug("Could not close realtime client websocket; it is already gone") - finally: - from litellm.litellm_core_utils.realtime_streaming import ( - REALTIME_SESSION_SUCCESS_LOGGED_KEY, - ) - - if not litellm_logging_obj.model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY): - await _release_realtime_budget_reservation(user_api_key_dict) + except Exception: # noqa: BLE001 # the lower layer may have closed the socket already; closing twice is not an error + verbose_proxy_logger.debug("Could not close realtime client websocket; it is already gone") finally: - await proxy_logging_obj._arelease_max_parallel_requests_on_disconnect(user_api_key_dict) # pyright: ignore[reportPrivateUsage] # same release idiom the HTTP disconnect path uses + from litellm.litellm_core_utils.realtime_streaming import ( + REALTIME_SESSION_SUCCESS_LOGGED_KEY, + ) + + if not litellm_logging_obj.model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY): + await _release_realtime_budget_reservation(user_api_key_dict) + await _release_realtime_max_parallel_slot(user_api_key_dict) ###################################################################### diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 11c2d301365..404deb3ca50 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -31,6 +31,7 @@ from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded from litellm.caching.dual_cache import DualCache from litellm.proxy._types import LitellmUserRoles, TokenCountRequest, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.hooks.parallel_request_limiter_v3 import RequestRateLimiterStash from litellm.proxy.proxy_server import app, initialize from litellm.utils import _invalidate_model_cost_lowercase_map @@ -10075,6 +10076,50 @@ async def test_successful_realtime_session_leaves_the_reservation_for_the_cost_c assert reservation["finalized"] is False +_LIT6463_COUNTER_KEY: Final = "{api_key:hashed-token}:max_parallel_requests" + + +async def _lit6463_drive_realtime_session_holding_a_max_parallel_slot( + *, + backend_logged_success: bool, + phase_one_exit: str | None = None, +) -> tuple[DualCache, RequestRateLimiterStash]: + """Run the realtime endpoint with a real v3 limiter registered and the request's + stash already holding slot-1 of a two-slot counter, the state pre-call leaves + behind. Returns the limiter's cache and the stash so the test can read what the + endpoint did to the slot.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + _PROXY_MaxParallelRequestsHandler_v3, + _request_stash, + ) + from litellm.proxy.utils import InternalUsageCache + + dual_cache: Final = DualCache() + await dual_cache.async_set_cache( + key=_LIT6463_COUNTER_KEY, value={"slot-1": 1.0, "slot-2": 2.0}, local_only=True + ) + limiter: Final = _PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(dual_cache)) + stash: Final = RequestRateLimiterStash( + parallel_slot={"slot_id": "slot-1", "counter_keys": [_LIT6463_COUNTER_KEY]} + ) + reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} + + stash_token: Final = _request_stash.set(stash) + try: + hooks: Final = patch.dict(ps.proxy_logging_obj.proxy_hook_mapping, {"parallel_request_limiter": limiter}) # test-quality-ok: registers a real limiter on the module-global hook map the route reads; assertion observes its counter + expected_exit: Final = ( + pytest.raises(asyncio.CancelledError) if phase_one_exit == "pre_call_cancelled" else contextlib.nullcontext() + ) + with hooks, expected_exit: + await _lit6973_drive_realtime_session( + reservation, backend_logged_success=backend_logged_success, phase_one_exit=phase_one_exit + ) + finally: + _request_stash.reset(stash_token) + return dual_cache, stash + + @pytest.mark.asyncio @pytest.mark.parametrize("phase_one_exit", [None, "pre_call", "pre_call_cancelled"]) async def test_realtime_session_ending_without_llm_callbacks_releases_the_max_parallel_slot( @@ -10086,39 +10131,33 @@ async def test_realtime_session_ending_without_llm_callbacks_releases_the_max_pa rejects the session, or the task is cancelled while still in pre-call) has to be released by the route itself, or the slot stays occupied until its TTL and the key's next session is refused with a 429.""" - from litellm.caching.caching import DualCache - from litellm.proxy import proxy_server as ps - from litellm.proxy.hooks.parallel_request_limiter_v3 import ( - RequestRateLimiterStash, - _PROXY_MaxParallelRequestsHandler_v3, - _request_stash, + dual_cache, stash = await _lit6463_drive_realtime_session_holding_a_max_parallel_slot( + backend_logged_success=False, phase_one_exit=phase_one_exit ) - from litellm.proxy.utils import InternalUsageCache - counter_key: Final = "{api_key:hashed-token}:max_parallel_requests" - dual_cache: Final = DualCache() - await dual_cache.async_set_cache(key=counter_key, value={"slot-1": 1.0, "slot-2": 2.0}, local_only=True) - limiter: Final = _PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(dual_cache)) - stash: Final = RequestRateLimiterStash(parallel_slot={"slot_id": "slot-1", "counter_keys": [counter_key]}) - reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} - - stash_token: Final = _request_stash.set(stash) - try: - hooks: Final = patch.dict(ps.proxy_logging_obj.proxy_hook_mapping, {"parallel_request_limiter": limiter}) # test-quality-ok: registers a real limiter on the module-global hook map the route reads; assertion observes its counter - expected_exit: Final = ( - pytest.raises(asyncio.CancelledError) if phase_one_exit == "pre_call_cancelled" else contextlib.nullcontext() - ) - with hooks, expected_exit: - await _lit6973_drive_realtime_session( - reservation, backend_logged_success=False, phase_one_exit=phase_one_exit - ) - finally: - _request_stash.reset(stash_token) - - assert await dual_cache.async_get_cache(key=counter_key, local_only=True) == {"slot-2": 2.0} + assert await dual_cache.async_get_cache(key=_LIT6463_COUNTER_KEY, local_only=True) == {"slot-2": 2.0} assert stash.parallel_slot is None +@pytest.mark.asyncio +async def test_successful_realtime_session_leaves_the_max_parallel_slot_for_the_limiter_callback(): + """A session that enqueued its success callback hands the slot to the limiter's + own success handler, which runs on the logging worker. If the route also released + it, the two releases would race on the same stashed acquisition and, under the + limiter's integer in-memory fallback, double-decrement the counter so the key + admits more sessions than max_parallel_requests allows. With the success stamp + present the route leaves the slot and the stash alone.""" + dual_cache, stash = await _lit6463_drive_realtime_session_holding_a_max_parallel_slot( + backend_logged_success=True + ) + + assert await dual_cache.async_get_cache(key=_LIT6463_COUNTER_KEY, local_only=True) == { + "slot-1": 1.0, + "slot-2": 2.0, + } + assert stash.parallel_slot == {"slot_id": "slot-1", "counter_keys": [_LIT6463_COUNTER_KEY]} + + @pytest.mark.asyncio async def test_release_or_invalidate_falls_back_to_invalidating_the_counters(): """If releasing the reservation itself fails (e.g. the counter store is down), From 336ead5106d6d4f7c9ad1d37ea2ab4b615664259 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 20:23:54 +0000 Subject: [PATCH 101/112] fix(proxy): import UserNotFoundError in grants resolver from the types module Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/resolvers/grants.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/auth/resolvers/grants.py b/litellm/proxy/auth/resolvers/grants.py index eb39d2a6812..cbfe21924ec 100644 --- a/litellm/proxy/auth/resolvers/grants.py +++ b/litellm/proxy/auth/resolvers/grants.py @@ -28,11 +28,11 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.auth_checks import ( TeamNotFoundError, - UserNotFoundError, get_team_membership, get_team_object, get_user_object, ) +from litellm.types.proxy.auth.auth_checks import UserNotFoundError if TYPE_CHECKING: from litellm.proxy._types import Span From 9a62a5ebee7743c93d179039aa8311b87a4cb33a Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 20:29:35 +0000 Subject: [PATCH 102/112] refactor(prometheus): source PROXY_LLM_PROVIDER_FALLBACK from litellm.constants Importing the fallback from litellm.proxy.hooks.rate_limiter_utils at the top of litellm/integrations/prometheus.py closed an import cycle CodeQL flagged on #41059 (rate_limiter_utils -> litellm -> ... -> prometheus). Hoist the constant into litellm/constants.py so both modules read it from a leaf module. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 3 +++ litellm/integrations/prometheus.py | 2 +- litellm/proxy/hooks/rate_limiter_utils.py | 3 +-- 3 files changed, 5 insertions(+), 3 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 5751e6e46af..aaa1e1484a2 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -227,6 +227,9 @@ PRE_CALL_EXECUTED_GUARDRAILS_KEY: Final = "_pre_call_executed_guardrails" # Attribute stamped on log_guardrail_information wrappers so __init_subclass__ does not wrap them again LOGS_GUARDRAIL_INFORMATION_MARKER: Final = "_litellm_logs_guardrail_information" +# llm_provider stamped on proxy-side rate limit errors when the model resolves to no deployment +PROXY_LLM_PROVIDER_FALLBACK: Final = "litellm_proxy" + # Generic fallback for unknown models DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET: Final = int( os.getenv("DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET", 128) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 2ba5d471e75..69a38e83835 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -16,6 +16,7 @@ from pydantic import BaseModel import litellm from litellm._logging import print_verbose, verbose_logger +from litellm.constants import PROXY_LLM_PROVIDER_FALLBACK from litellm.exceptions import ( validate_rate_limit_category, validate_rate_limit_type, @@ -41,7 +42,6 @@ from litellm.proxy._types import ( LiteLLM_UserTable, UserAPIKeyAuth, ) -from litellm.proxy.hooks.rate_limiter_utils import PROXY_LLM_PROVIDER_FALLBACK from litellm.repositories.base_repository import BaseRepository from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.team_repository import TeamRepository diff --git a/litellm/proxy/hooks/rate_limiter_utils.py b/litellm/proxy/hooks/rate_limiter_utils.py index cdc3896c10d..bbe38068638 100644 --- a/litellm/proxy/hooks/rate_limiter_utils.py +++ b/litellm/proxy/hooks/rate_limiter_utils.py @@ -6,11 +6,10 @@ from typing import Final import litellm from litellm._logging import verbose_proxy_logger +from litellm.constants import PROXY_LLM_PROVIDER_FALLBACK from litellm.types.router import ModelGroupInfo from litellm.types.utils import PriorityReservationDict -PROXY_LLM_PROVIDER_FALLBACK: Final = "litellm_proxy" - def resolve_llm_provider_for_rate_limit( model: str | None, From c8440b5638881d9eb463ecbc5a17ba4c7351ee3b Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 20:34:58 +0000 Subject: [PATCH 103/112] fix(proxy): move credentials hint helper into a leaf html_forms module to clear CodeQL cyclic import Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../html_forms/default_credentials_hint.py | 10 ++++++++++ .../discovery_endpoints/ui_discovery_endpoints.py | 10 +--------- litellm/proxy/management_endpoints/ui_sso.py | 2 +- litellm/proxy/proxy_server.py | 2 +- 4 files changed, 13 insertions(+), 11 deletions(-) create mode 100644 litellm/proxy/common_utils/html_forms/default_credentials_hint.py diff --git a/litellm/proxy/common_utils/html_forms/default_credentials_hint.py b/litellm/proxy/common_utils/html_forms/default_credentials_hint.py new file mode 100644 index 00000000000..e3ed948501b --- /dev/null +++ b/litellm/proxy/common_utils/html_forms/default_credentials_hint.py @@ -0,0 +1,10 @@ +import os +from collections.abc import Mapping + + +def should_hide_default_credentials_hint(general_settings: Mapping[str, object]) -> bool: + return ( + os.getenv("LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT", "false").lower() == "true" + or general_settings.get("hide_default_credentials_hint", False) is True + or bool(os.getenv("UI_PASSWORD")) + ) diff --git a/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py b/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py index e0efe9dea2c..8b042d18cd0 100644 --- a/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py +++ b/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py @@ -1,10 +1,10 @@ #### Analytics Endpoints ##### import os -from collections.abc import Mapping from typing import Final from fastapi import APIRouter +from litellm.proxy.common_utils.html_forms.default_credentials_hint import should_hide_default_credentials_hint from litellm.types.proxy.discovery_endpoints.ui_discovery_endpoints import ( UiDiscoveryEndpoints, ) @@ -12,14 +12,6 @@ from litellm.types.proxy.discovery_endpoints.ui_discovery_endpoints import ( router: Final = APIRouter() -def should_hide_default_credentials_hint(general_settings: Mapping[str, object]) -> bool: - return ( - os.getenv("LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT", "false").lower() == "true" - or general_settings.get("hide_default_credentials_hint", False) is True - or bool(os.getenv("UI_PASSWORD")) - ) - - @router.get("/.well-known/litellm-ui-config", response_model=UiDiscoveryEndpoints) @router.get("/litellm/.well-known/litellm-ui-config", response_model=UiDiscoveryEndpoints) # if mounted at root path async def get_ui_config(): diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 6a775998f4f..091dccf1433 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -101,12 +101,12 @@ from litellm.proxy.common_utils.admin_ui_utils import ( admin_ui_disabled, show_missing_vars_in_env, ) +from litellm.proxy.common_utils.html_forms.default_credentials_hint import should_hide_default_credentials_hint from litellm.proxy.common_utils.html_forms.jwt_display_template import ( jwt_display_template, ) from litellm.proxy.common_utils.html_forms.ui_login import build_ui_login_form from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache -from litellm.proxy.discovery_endpoints.ui_discovery_endpoints import should_hide_default_credentials_hint from litellm.proxy.management_endpoints.internal_user_endpoints import new_user from litellm.proxy.management_endpoints.sso import CustomMicrosoftSSO from litellm.proxy.management_endpoints.sso.id_jag_assertion_capture import ( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 230513d7eb4..fce468aa564 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -365,6 +365,7 @@ from litellm.proxy.common_utils.healthy_model_filter import ( get_hidden_unhealthy_model_names, is_healthy_only_listing_default, ) +from litellm.proxy.common_utils.html_forms.default_credentials_hint import should_hide_default_credentials_hint from litellm.proxy.common_utils.html_forms.ui_login import build_ui_login_form from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, @@ -450,7 +451,6 @@ from litellm.proxy.discovery_endpoints import ( agent_skills_discovery_router, ui_discovery_endpoints_router, ) -from litellm.proxy.discovery_endpoints.ui_discovery_endpoints import should_hide_default_credentials_hint from litellm.proxy.fine_tuning_endpoints.endpoints import router as fine_tuning_router from litellm.proxy.fine_tuning_endpoints.endpoints import set_fine_tuning_config from litellm.proxy.google_endpoints.endpoints import router as google_router From 3177c37e22c72b58d5507265c0b9205bbd61b828 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Mon, 14 Sep 2026 12:57:59 -0700 Subject: [PATCH 104/112] fix(cli): show routed models and session stats for LLM API keys --- litellm/proxy/auth/route_checks.py | 3 ++ litellm/proxy/client/cli/README.md | 4 +- .../client/cli/commands/statusline_script.py | 25 +++++----- .../proxy/auth/test_route_checks.py | 47 +++++++++++++++++-- .../client/cli/test_statusline_script.py | 28 ++++++++--- 5 files changed, 83 insertions(+), 24 deletions(-) diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index 953e3cf3e88..1789b080897 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -136,6 +136,9 @@ class RouteChecks: # For llm_api_routes, also check registered pass-through endpoints ################################################ if allowed_route == "llm_api_routes": + if route == "/auto_router/session" and RouteChecks._get_request_method(request) == "GET": + return True + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( InitPassThroughEndpointHelpers, ) diff --git a/litellm/proxy/client/cli/README.md b/litellm/proxy/client/cli/README.md index 3b0ff9d7add..cb867cf9e61 100644 --- a/litellm/proxy/client/cli/README.md +++ b/litellm/proxy/client/cli/README.md @@ -585,7 +585,9 @@ LiteLLM ████████░░░░░░░░░░░░░░ Claude Opus 5 ████████████████████████ $0.38 ``` -The routed model comes from Claude Code's own transcript, so it only names the tier model when the auto-router deployment sets `return_raw_model_name: true` (the `lite autoroute` wizard does); otherwise it shows the alias you requested. The cost lines come from `GET /auto_router/session?session_id=...`, which any virtual key may call for its own sessions, and are cached for five seconds under a per-user `$TMPDIR/litellm-statusline-` directory. The baseline is the priciest model in the router's hardest tier, the same counterfactual the auto-router's savings reports use. `lite unconfigure claude` removes the `statusLine` entry only while it still points at that script. +After the first response, the status line uses the latest routed model recorded by `GET /auto_router/session?session_id=...`, so it can show the tier model even when the transcript contains the router alias. If no session record is available, it falls back to Claude Code's transcript. Session records and costs are cached for five seconds under a per-user `$TMPDIR/litellm-statusline-` directory. The gateway records turns asynchronously, so the display can briefly lag a completed turn. Any virtual key may read its own sessions. The baseline is the priciest model in the router's hardest tier, the same counterfactual the auto-router's savings reports use. `lite unconfigure claude` removes the `statusLine` entry only while it still points at that script + +After upgrading the CLI, rerun your original `lite configure claude` command with the same gateway, key and model choice to refresh `~/.litellm/statusline.py`. Keep any explicit `--model` value: omitting it removes the earlier model pin. Package upgrades alone do not refresh this installed copy `lite codex` registers the same script as a Codex `Stop` hook for the launch, so after each turn Codex prints the same block as a system message. Codex asks once to trust the hook; the answer is remembered for later launches. diff --git a/litellm/proxy/client/cli/commands/statusline_script.py b/litellm/proxy/client/cli/commands/statusline_script.py index a8abeb68978..5493586f627 100644 --- a/litellm/proxy/client/cli/commands/statusline_script.py +++ b/litellm/proxy/client/cli/commands/statusline_script.py @@ -7,19 +7,17 @@ status refresh (about every 300ms while typing), so the proxy is asked at most o TTL per session and every other refresh is served from a small on-disk cache that holds only the proxy's answer, never the key. -Claude Code pipes a JSON payload on stdin (session_id, transcript_path, model); the routed -model is the `message.model` of the latest foreground assistant line in the transcript, -which is the proxy's response `model` field. That only names the tier model when the -auto-router deployment sets `return_raw_model_name: true`; otherwise it is the alias the -client requested. Codex pipes its Stop event instead (hook_event_name, session_id) and has -no transcript to read, so the routed model comes from the proxy's session record and the -result is printed as a `systemMessage` for the transcript. The proxy key is read from the -agent's own environment (the static token `lite configure claude` writes); nothing here -spawns a credential helper. +Claude Code pipes a JSON payload on stdin (session_id, transcript_path, model). After the +first foreground assistant response, the routed model comes from the proxy's session +record, falling back to the latest foreground assistant `message.model` in the transcript +when no record is available. Codex pipes its Stop event instead (hook_event_name, session_id) +and prints the session record as a `systemMessage` for the transcript. The proxy key is read +from the agent's own environment (the static token `lite configure claude` writes); nothing +here spawns a credential helper. -Cost figures come from GET /auto_router/session on the proxy, which reads the per-session -rollup written by the spend flush. That flush is asynchronous, so a turn's cost lands a -second or two after the turn; the cache TTL absorbs it. +The routed model and cost figures come from GET /auto_router/session on the proxy, which +reads the per-session rollup written by the asynchronous spend flush. The record and cache +can briefly lag a completed turn. """ from __future__ import annotations @@ -348,7 +346,8 @@ def status_line( if not session_id or not credentials.usable: return render(label, None, config_dir, color_enabled(env)) session: Final = load_session(credentials, session_id, cache_dir, fetch) - return render(label, session, config_dir, color_enabled(env)) + routed_label: Final = model_label(session.last_model, config_dir) if session is not None else label + return render(routed_label, session, config_dir, color_enabled(env)) def codex_stop_message( diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index c8b3d789665..0950b56bf03 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -1,5 +1,6 @@ import os from datetime import datetime +from typing import Final from unittest.mock import MagicMock, patch @@ -3919,10 +3920,15 @@ def test_claude_code_marketplace_routes_open_to_internal_users(route): @pytest.mark.parametrize("user_role", [None, LitellmUserRoles.INTERNAL_USER.value, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value]) -def test_auto_router_session_is_reachable_by_any_key_but_benchmarks_stays_admin_only(user_role): - valid_token = UserAPIKeyAuth(api_key="hash-of-caller", user_role=user_role) - request = MagicMock(spec=Request) - request.query_params = {"session_id": "sess-1"} +@pytest.mark.parametrize("allowed_routes", [None, ["llm_api_routes"]]) +def test_auto_router_session_is_reachable_by_any_key_but_benchmarks_stays_admin_only( + user_role: str | None, allowed_routes: list[str] | None +) -> None: + valid_token: Final = UserAPIKeyAuth(api_key="hash-of-caller", user_role=user_role, allowed_routes=allowed_routes) + request: Final = Request({"type": "http", "method": "GET", "query_string": b"session_id=sess-1"}) + + assert RouteChecks.should_call_route("/auto_router/session", valid_token, request) is True + assert RouteChecks.is_llm_api_route("/auto_router/session") is False RouteChecks.non_proxy_admin_allowed_routes_check( user_obj=None, @@ -3941,3 +3947,36 @@ def test_auto_router_session_is_reachable_by_any_key_but_benchmarks_stays_admin_ valid_token=valid_token, request_data={}, ) + + +@pytest.mark.parametrize( + "route,method,allowed_routes", + [ + ("/auto_router/session", method, ["llm_api_routes"]) + for method in ("POST", "PUT", "PATCH", "DELETE", "HEAD", "OPTIONS", None) + ] + + [ + (route, "GET", ["llm_api_routes"]) + for route in ( + "/auto_router/benchmarks", + "/auto_router/test_routing", + "/auto_router/validate_complexity_router_config", + "/auto_router/session/other", + "/auto_router/sessions", + ) + ] + + [ + ("/auto_router/session", "GET", allowed_routes) + for allowed_routes in (["/v1/messages"], ["info_routes"], ["openai_routes"]) + ], +) +def test_auto_router_session_read_grant_rejects_other_methods_paths_and_scopes( + route: str, method: str | None, allowed_routes: list[str] +) -> None: + valid_token: Final = UserAPIKeyAuth(api_key="hash-of-caller", allowed_routes=allowed_routes) + request: Final = Request({"type": "http", "method": method}) if method is not None else None + + with pytest.raises(HTTPException) as error: + RouteChecks.should_call_route(route, valid_token, request) + + assert error.value.status_code == 403 diff --git a/tests/test_litellm/proxy/client/cli/test_statusline_script.py b/tests/test_litellm/proxy/client/cli/test_statusline_script.py index 691764fbef4..5c0cf6b5703 100644 --- a/tests/test_litellm/proxy/client/cli/test_statusline_script.py +++ b/tests/test_litellm/proxy/client/cli/test_statusline_script.py @@ -8,6 +8,7 @@ import re import subprocess import sys from pathlib import Path +from typing import Final import pytest @@ -297,16 +298,31 @@ class TestRender: class TestClaudeCodeMode: - def test_the_transcript_names_the_routed_model_and_the_proxy_adds_the_savings(self, tmp_path, transcript, config_dir): - seen = [] + @pytest.mark.parametrize("transcript_model", ("claude-auto", "anthropic/claude-opus-5")) + def test_the_session_names_the_routed_model_even_when_the_transcript_differs( + self, tmp_path: Path, config_dir: Path, transcript_model: str + ) -> None: + transcript: Final = tmp_path / "session.jsonl" + transcript.write_text(_assistant_line(transcript_model) + "\n") - def fetch(credentials, session_id): - seen.append((credentials, session_id)) + def fetch(credentials: Credentials, session_id: str) -> Fetched: + assert credentials == Credentials("http://127.0.0.1:4000", "sk-virtual") + assert session_id == SESSION_ID return Fetched(RECORDED, definitive=True) - text = _run(_payload(transcript), _env(tmp_path, config_dir), fetch) + text: Final = _run(_payload(transcript), _env(tmp_path, config_dir), fetch) assert text.startswith("claude-auto · Routed to: claude-sonnet-5 -63% vs Claude Opus 5\n") - assert seen == [(Credentials("http://127.0.0.1:4000", "sk-virtual"), SESSION_ID)] + + def test_a_discovered_display_name_labels_the_sessions_model( + self, tmp_path: Path, transcript: Path, config_dir: Path + ) -> None: + session: Final = RECORDED._replace(last_model="anthropic/claude-opus-5") + + def fetch(credentials: Credentials, session_id: str) -> Fetched: + return Fetched(session, definitive=True) + + text: Final = _run(_payload(transcript), _env(tmp_path, config_dir), fetch) + assert text.startswith("claude-auto · Routed to: Claude Opus 5 -63% vs Claude Opus 5\n") def test_an_unrecorded_session_degrades_to_the_routed_line(self, tmp_path, transcript, config_dir): assert _run(_payload(transcript), _env(tmp_path, config_dir), lambda c, s: Fetched(None, True)) == ( From 91c964a338c759cb5ea69403d90a12696795d816 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 20:46:09 +0000 Subject: [PATCH 105/112] fix(auth): evict the membership cache entry when invalidation lands during the cache write Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/auth_checks.py | 15 +++++---- .../proxy/auth/test_auth_checks.py | 32 +++++++++++++++++++ 2 files changed, 40 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 43d840130b6..22644bcd207 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2212,13 +2212,14 @@ async def _fetch_team_membership_from_db( value=NO_TEAM_MEMBERSHIP_SENTINEL, ttl=get_management_object_ttl(user_api_key_cache), ) - return None - - await user_api_key_cache.async_set_cache( - key=_key, - value=membership, - model_type=LiteLLM_TeamMembership, - ) + else: + await user_api_key_cache.async_set_cache( + key=_key, + value=membership, + model_type=LiteLLM_TeamMembership, + ) + if _team_membership_inflight.get(_key) is not asyncio.current_task(): + await user_api_key_cache.async_delete_cache(key=_key) return membership diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 986d53b80b0..61f0675aa1e 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -6543,6 +6543,38 @@ async def test_get_team_membership_invalidation_mid_flight_discards_stale_load() assert cached is not None and cached.budget_id == "budget-new" +@pytest.mark.asyncio +async def test_get_team_membership_invalidation_during_cache_write_evicts_stale_entry(): + from litellm.proxy.auth.auth_checks import get_team_membership, invalidate_team_member_spend_state + from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key + + write_started = asyncio.Event() + release_write = asyncio.Event() + + class _SlowWriteCache(UserApiKeyCache): + async def async_set_cache(self, key, value, local_only=False, **kwargs): + write_started.set() + await release_write.wait() + return await super().async_set_cache(key, value, local_only=local_only, **kwargs) + + row = MagicMock() + row.dict = lambda: {"user_id": "u-w", "team_id": "t-w", "spend": 1.0, "budget_id": "budget-old"} + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=row) + cache = _SlowWriteCache() + + stale = asyncio.create_task( + get_team_membership(user_id="u-w", team_id="t-w", prisma_client=mock_prisma_client, user_api_key_cache=cache) + ) + await asyncio.wait_for(write_started.wait(), timeout=2) + await invalidate_team_member_spend_state(user_id="u-w", team_id="t-w", user_api_key_cache=cache) + release_write.set() + stale_result = await stale + + assert stale_result is not None and stale_result.budget_id == "budget-old" + assert await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-w", team_id="t-w")) is None + + @pytest.mark.asyncio async def test_common_checks_calls_get_team_membership_once_per_request(): from fastapi import Request From abc85ba6078b57466fd7d7a0d20217c578844add Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 20:47:37 +0000 Subject: [PATCH 106/112] fix(proxy): leave the realtime max_parallel_requests slot to the limiter failure callback when a refusal was logged Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_core_utils/realtime_streaming.py | 2 + litellm/proxy/proxy_server.py | 9 ++++- .../test_realtime_streaming.py | 21 ++++++++++ tests/test_litellm/proxy/test_proxy_server.py | 40 +++++++++++++++++-- 4 files changed, 66 insertions(+), 6 deletions(-) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index e3f8786a39a..2177c999804 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -36,6 +36,7 @@ else: REALTIME_SESSION_SUCCESS_LOGGED_KEY: Final = "realtime_session_success_logged" +REALTIME_SESSION_FAILURE_LOGGED_KEY: Final = "realtime_session_failure_logged" @dataclass(frozen=True, slots=True) @@ -1153,6 +1154,7 @@ class RealTimeStreaming: self._logging_worker.ensure_initialized_and_enqueue( self.logging_obj.dispatch_failure_handlers(error, traceback.format_exc(), prefer_async_handlers=True) ) + self.logging_obj.model_call_details[REALTIME_SESSION_FAILURE_LOGGED_KEY] = True @staticmethod def _detect_beta_header(websocket: ScopedWebSocket) -> bool: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f37744f4685..74783398456 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11894,7 +11894,10 @@ async def _release_realtime_budget_reservation(user_api_key_dict: UserAPIKeyAuth async def _release_realtime_max_parallel_slot(user_api_key_dict: UserAPIKeyAuth) -> None: - await proxy_logging_obj._arelease_max_parallel_requests_on_disconnect(user_api_key_dict) # pyright: ignore[reportPrivateUsage] # same release idiom the HTTP disconnect path uses + release_like_http_disconnect: Final = ( + proxy_logging_obj._arelease_max_parallel_requests_on_disconnect # pyright: ignore[reportPrivateUsage] # shared + ) + await release_like_http_disconnect(user_api_key_dict) async def _reject_realtime_session( @@ -12053,12 +12056,14 @@ async def realtime_websocket_endpoint( verbose_proxy_logger.debug("Could not close realtime client websocket; it is already gone") finally: from litellm.litellm_core_utils.realtime_streaming import ( + REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY, ) if not litellm_logging_obj.model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY): await _release_realtime_budget_reservation(user_api_key_dict) - await _release_realtime_max_parallel_slot(user_api_key_dict) + if not litellm_logging_obj.model_call_details.get(REALTIME_SESSION_FAILURE_LOGGED_KEY): + await _release_realtime_max_parallel_slot(user_api_key_dict) ###################################################################### diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py index 2a33d84ec78..e1eb61b59b9 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -12,6 +12,7 @@ from websockets.frames import Close import litellm from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.realtime_streaming import ( + REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY, RealTimeStreaming, client_sent_openai_beta_realtime_header, @@ -3399,6 +3400,26 @@ async def test_refused_session_does_not_stamp_the_reservation_ownership_marker() assert REALTIME_SESSION_SUCCESS_LOGGED_KEY not in session.logging.model_call_details +@pytest.mark.asyncio +async def test_refused_session_stamps_the_failure_ownership_marker(): + """LIT-6463: the enqueued failure callback releases the key's max_parallel_requests + slot from the logging worker, so a refusal stamps REALTIME_SESSION_FAILURE_LOGGED_KEY. + The proxy endpoint reads it to leave the slot to that callback instead of racing it. + A session that relayed frames logs a success and must not carry the failure stamp.""" + upstream_close: Final = ConnectionClosed(Close(1008, _UPSTREAM_REFUSAL), None) + refused: Final = _relay_session(_client_ws_that_never_sends(), _backend_ws_closing_with(upstream_close)) + session_created: Final = json.dumps({"type": "session.created", "session": {"id": "sess_1"}}).encode() + relayed: Final = _relay_session( + _client_ws_that_never_sends(), _backend_ws_closing_with(session_created, upstream_close) + ) + + await refused.run() + await relayed.run() + + assert refused.logging.model_call_details.get(REALTIME_SESSION_FAILURE_LOGGED_KEY) is True + assert REALTIME_SESSION_FAILURE_LOGGED_KEY not in relayed.logging.model_call_details + + @pytest.mark.asyncio async def test_transformed_transcription_completion_never_sends_response_create(): from typing import Final diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 404deb3ca50..986cac2a18e 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -9913,6 +9913,7 @@ async def _lit6973_drive_realtime_session( reservation: dict, *, backend_logged_success: bool, + backend_logged_failure: bool = False, phase_one_exit: str | None = None, websocket: MagicMock | None = None, ) -> MagicMock: @@ -9932,7 +9933,10 @@ async def _lit6973_drive_realtime_session( logging object carries a real model_call_details dict so the stamp is observable, and the reservation has empty entries so the real release touches no counter store.""" - from litellm.litellm_core_utils.realtime_streaming import REALTIME_SESSION_SUCCESS_LOGGED_KEY + from litellm.litellm_core_utils.realtime_streaming import ( + REALTIME_SESSION_FAILURE_LOGGED_KEY, + REALTIME_SESSION_SUCCESS_LOGGED_KEY, + ) from litellm.proxy import proxy_server as ps user_api_key_dict: Final = UserAPIKeyAuth(api_key="sk-test", token="hashed-token") @@ -9944,6 +9948,8 @@ async def _lit6973_drive_realtime_session( async def fake_llm_call() -> None: if backend_logged_success: logging_obj.model_call_details[REALTIME_SESSION_SUCCESS_LOGGED_KEY] = True + if backend_logged_failure: + logging_obj.model_call_details[REALTIME_SESSION_FAILURE_LOGGED_KEY] = True from litellm.proxy._types import ProxyException @@ -10082,6 +10088,7 @@ _LIT6463_COUNTER_KEY: Final = "{api_key:hashed-token}:max_parallel_requests" async def _lit6463_drive_realtime_session_holding_a_max_parallel_slot( *, backend_logged_success: bool, + backend_logged_failure: bool = False, phase_one_exit: str | None = None, ) -> tuple[DualCache, RequestRateLimiterStash]: """Run the realtime endpoint with a real v3 limiter registered and the request's @@ -10107,13 +10114,20 @@ async def _lit6463_drive_realtime_session_holding_a_max_parallel_slot( stash_token: Final = _request_stash.set(stash) try: - hooks: Final = patch.dict(ps.proxy_logging_obj.proxy_hook_mapping, {"parallel_request_limiter": limiter}) # test-quality-ok: registers a real limiter on the module-global hook map the route reads; assertion observes its counter + hooks: Final = patch.dict( # test-quality-ok: registers the real limiter the route's release reads + ps.proxy_logging_obj.proxy_hook_mapping, {"parallel_request_limiter": limiter} + ) expected_exit: Final = ( - pytest.raises(asyncio.CancelledError) if phase_one_exit == "pre_call_cancelled" else contextlib.nullcontext() + pytest.raises(asyncio.CancelledError) + if phase_one_exit == "pre_call_cancelled" + else contextlib.nullcontext() ) with hooks, expected_exit: await _lit6973_drive_realtime_session( - reservation, backend_logged_success=backend_logged_success, phase_one_exit=phase_one_exit + reservation, + backend_logged_success=backend_logged_success, + backend_logged_failure=backend_logged_failure, + phase_one_exit=phase_one_exit, ) finally: _request_stash.reset(stash_token) @@ -10158,6 +10172,24 @@ async def test_successful_realtime_session_leaves_the_max_parallel_slot_for_the_ assert stash.parallel_slot == {"slot_id": "slot-1", "counter_keys": [_LIT6463_COUNTER_KEY]} +@pytest.mark.asyncio +async def test_refused_realtime_session_leaves_the_max_parallel_slot_for_the_limiter_failure_callback(): + """An upstream refusal before any frame enqueues the failure callback instead, and + the limiter's failure handler releases the slot from the logging worker just like + the success handler does. The route sees no success stamp, so it still settles the + budget reservation, but it must leave the slot to that callback or the two releases + race on the same acquisition.""" + dual_cache, stash = await _lit6463_drive_realtime_session_holding_a_max_parallel_slot( + backend_logged_success=False, backend_logged_failure=True + ) + + assert await dual_cache.async_get_cache(key=_LIT6463_COUNTER_KEY, local_only=True) == { + "slot-1": 1.0, + "slot-2": 2.0, + } + assert stash.parallel_slot == {"slot_id": "slot-1", "counter_keys": [_LIT6463_COUNTER_KEY]} + + @pytest.mark.asyncio async def test_release_or_invalidate_falls_back_to_invalidating_the_counters(): """If releasing the reservation itself fails (e.g. the counter store is down), From ffeea30f23d4de5dbd7554fa3a0518ab727e168a Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 20:55:22 +0000 Subject: [PATCH 107/112] test(auth): cover a stale membership write landing after a fresh reload Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/auth/test_auth_checks.py | 45 +++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 61f0675aa1e..96e9958a55d 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -6575,6 +6575,51 @@ async def test_get_team_membership_invalidation_during_cache_write_evicts_stale_ assert await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-w", team_id="t-w")) is None +@pytest.mark.asyncio +async def test_get_team_membership_stale_write_finishing_after_fresh_load_never_serves_old_row(): + from litellm.proxy.auth.auth_checks import get_team_membership, invalidate_team_member_spend_state + + write_started = asyncio.Event() + release_stale_write = asyncio.Event() + stale_writes = iter((release_stale_write,)) + + class _SlowFirstWriteCache(UserApiKeyCache): + async def async_set_cache(self, key, value, local_only=False, **kwargs): + release = next(stale_writes, None) + if release is not None: + write_started.set() + await release.wait() + return await super().async_set_cache(key, value, local_only=local_only, **kwargs) + + rows = iter(("budget-old", "budget-new", "budget-new")) + + async def _find_unique(*args, **kwargs): + row = MagicMock() + row.dict = lambda: {"user_id": "u-sw", "team_id": "t-sw", "spend": 1.0, "budget_id": next(rows)} + return row + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(side_effect=_find_unique) + cache = _SlowFirstWriteCache() + + async def _load(): + return await get_team_membership( + user_id="u-sw", team_id="t-sw", prisma_client=mock_prisma_client, user_api_key_cache=cache + ) + + stale = asyncio.create_task(_load()) + await asyncio.wait_for(write_started.wait(), timeout=2) + await invalidate_team_member_spend_state(user_id="u-sw", team_id="t-sw", user_api_key_cache=cache) + fresh_result = await _load() + release_stale_write.set() + stale_result = await stale + after_result = await _load() + + assert fresh_result is not None and fresh_result.budget_id == "budget-new" + assert stale_result is not None and stale_result.budget_id == "budget-old" + assert after_result is not None and after_result.budget_id == "budget-new" + + @pytest.mark.asyncio async def test_common_checks_calls_get_team_membership_once_per_request(): from fastapi import Request From da839d4a117ed15e21d77aa3e8912271f3fe9c6c Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 14 Sep 2026 14:01:28 -0700 Subject: [PATCH 108/112] fix(harness): skip unavailable Rust traces --- .../strategies/trace_parity/fixtures.py | 13 ++++--- .../gateway/chat_completions/case.py | 4 +- .../trace_parity/gateway/execution.py | 7 ++-- .../trace_parity/gateway/responses/case.py | 3 +- .../strategies/trace_parity/models.py | 3 +- .../strategies/trace_parity/sdk/execution.py | 13 +++++-- .../trace_parity/sdk/responses/case.py | 2 +- .../strategies/trace_parity/test_runner.py | 38 +++++++++++++++++++ 8 files changed, 63 insertions(+), 20 deletions(-) diff --git a/tests/rust-python-harness/strategies/trace_parity/fixtures.py b/tests/rust-python-harness/strategies/trace_parity/fixtures.py index 84ac525efb6..a3d084c97de 100644 --- a/tests/rust-python-harness/strategies/trace_parity/fixtures.py +++ b/tests/rust-python-harness/strategies/trace_parity/fixtures.py @@ -70,16 +70,17 @@ def aws_event_stream_frame(payload: Mapping[str, object]) -> bytes: def aws_event_stream_response( events: Iterable[Mapping[str, object]], *, corrupt_last_frame: bool = False ) -> RecordedHttpStreamResponse: - frames: Final = [aws_event_stream_frame(event) for event in events] - if corrupt_last_frame: - corrupted: Final = bytearray(frames[-1]) - corrupted[-1] ^= 0xFF - frames[-1] = bytes(corrupted) + frames: Final = tuple(aws_event_stream_frame(event) for event in events) + body: Final = ( + b"".join((*frames[:-1], frames[-1][:-1] + bytes((frames[-1][-1] ^ 0xFF,)))) + if corrupt_last_frame + else b"".join(frames) + ) return RecordedHttpStreamResponse( kind="http_stream", status_code=200, headers=AWS_EVENT_STREAM_HEADERS, - chunks=(RecordedStreamChunk.from_bytes(b"".join(frames)),), + chunks=(RecordedStreamChunk.from_bytes(body),), ) diff --git a/tests/rust-python-harness/strategies/trace_parity/gateway/chat_completions/case.py b/tests/rust-python-harness/strategies/trace_parity/gateway/chat_completions/case.py index 522d2d0259c..3dc6d731b4b 100644 --- a/tests/rust-python-harness/strategies/trace_parity/gateway/chat_completions/case.py +++ b/tests/rust-python-harness/strategies/trace_parity/gateway/chat_completions/case.py @@ -9,8 +9,6 @@ from ...models import GatewayRouteSpec, RouteFixture, TraceScenario, TraceSuite MAPPINGS: Final = ( mapping(span="python_chat_gateway_route", python_frame=r"proxy_server\.py:\d+ chat_completion$"), mapping(span="python_gateway_service", python_frame=r"ProxyBaseLLMRequestProcessing\.base_process_llm_request$"), - mapping(rust_span="chat_completions_gateway_route"), - mapping(rust_span="chat_completions"), mapping(span="python_chat_entrypoint", python_frame=r"main\.py:\d+ a?completion$"), mapping(span="python_provider_config", python_frame=r"ProviderConfigManager\.get_provider_chat_config$"), mapping(rust_span="validate_environment", python_frame=r"(? RouteFixture: TRACE_SUITE: Final = TraceSuite( - route=GatewayRouteSpec("chat_completions"), + route=GatewayRouteSpec("chat_completions", rust_supported=False), scenarios=( TraceScenario(name="async-anthropic", fixture=_fixture, mappings=MAPPINGS, asynchronous=True), TraceScenario( diff --git a/tests/rust-python-harness/strategies/trace_parity/gateway/execution.py b/tests/rust-python-harness/strategies/trace_parity/gateway/execution.py index 5f24760dfe3..94d6be7cebf 100644 --- a/tests/rust-python-harness/strategies/trace_parity/gateway/execution.py +++ b/tests/rust-python-harness/strategies/trace_parity/gateway/execution.py @@ -176,8 +176,9 @@ def execute_gateway_trace( scenario: TraceScenario, engine: TraceEngine = "both", ) -> TraceArtifact: - python_trace: Final = _collect(route, scenario, "python") if engine != "rust" else () - rust_trace: Final = _collect(route, scenario, "rust") if engine != "python" else () + effective_engine: Final[TraceEngine] = "python" if engine == "both" and not route.rust_supported else engine + python_trace: Final = _collect(route, scenario, "python") if effective_engine != "rust" else () + rust_trace: Final = _collect(route, scenario, "rust") if effective_engine != "python" else () collection_python_error: Final = None if isinstance(python_trace, tuple) else f"python: {python_trace.message}" rust_error: Final = None if isinstance(rust_trace, tuple) else f"rust: {rust_trace.message}" python_events: Final = python_trace if isinstance(python_trace, tuple) else () @@ -185,7 +186,7 @@ def execute_gateway_trace( python, rust, projection_error = _projections(python_events, rust_events) python_error: Final = projection_error or collection_python_error return TraceArtifact.from_traces( - engine=engine, + engine=effective_engine, surface="gateway", sdk_function=route.route, scenario=scenario.name, diff --git a/tests/rust-python-harness/strategies/trace_parity/gateway/responses/case.py b/tests/rust-python-harness/strategies/trace_parity/gateway/responses/case.py index 52c1ebb391f..7d805a0603c 100644 --- a/tests/rust-python-harness/strategies/trace_parity/gateway/responses/case.py +++ b/tests/rust-python-harness/strategies/trace_parity/gateway/responses/case.py @@ -11,7 +11,6 @@ MAPPINGS: Final = ( span="python_responses_gateway_route", python_frame=r"response_api_endpoints/endpoints\.py:\d+ responses_api$" ), mapping(span="python_gateway_service", python_frame=r"ProxyBaseLLMRequestProcessing\.base_process_llm_request$"), - mapping(rust_span="responses_gateway_route"), mapping(span="python_responses", python_frame=r"responses/main\.py:\d+ a?responses$"), mapping(span="python_provider_config", python_frame=r"ProviderConfigManager\.get_provider_responses_api_config$"), mapping(rust_span="validate_environment", python_frame=r"OpenAIResponsesAPIConfig\.validate_environment$"), @@ -50,7 +49,7 @@ def _stream_fixture(engine: Engine, base_url: str) -> RouteFixture: TRACE_SUITE: Final = TraceSuite( - route=GatewayRouteSpec("responses"), + route=GatewayRouteSpec("responses", rust_supported=False), scenarios=( TraceScenario(name="async-openai", fixture=_fixture, mappings=MAPPINGS, asynchronous=True), TraceScenario( diff --git a/tests/rust-python-harness/strategies/trace_parity/models.py b/tests/rust-python-harness/strategies/trace_parity/models.py index 3701fb41e14..d6ed42250c4 100644 --- a/tests/rust-python-harness/strategies/trace_parity/models.py +++ b/tests/rust-python-harness/strategies/trace_parity/models.py @@ -48,13 +48,14 @@ class RouteFixture: class RouteSpec: route: SdkFunction python_entrypoints: tuple[str, str] - rust_entrypoints: tuple[str, str] + rust_entrypoints: tuple[str, str] | None fixture: Callable[[Engine, str], RouteFixture] @dataclass(frozen=True, slots=True) class GatewayRouteSpec: route: SdkFunction + rust_supported: bool = True TraceRouteSpec: TypeAlias = RouteSpec | GatewayRouteSpec diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py b/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py index ffdbe842e85..783c22a0dc0 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py @@ -62,6 +62,8 @@ def _entrypoint(spec: RouteSpec, engine: Engine, *, asynchronous: bool) -> SdkCa from litellm.rust_bridge import get_native_bridge if engine == "rust": + if spec.rust_entrypoints is None: + return TraceExecutionFailure("rust", f"{spec.route} has no native Rust trace entrypoint") bridge: Final = cast(object | None, get_native_bridge()) if bridge is None: return TraceExecutionFailure("rust", "native Rust bridge is required for trace parity") @@ -153,6 +155,7 @@ def execute_trace( surface: Surface, engine: TraceEngine = "both", ) -> TraceArtifact: + effective_engine: Final[TraceEngine] = "python" if engine == "both" and route.rust_entrypoints is None else engine scenario_route: Final = RouteSpec( route=route.route, python_entrypoints=route.python_entrypoints, @@ -165,11 +168,13 @@ def execute_trace( "python", asynchronous=scenario.asynchronous, ) - if engine != "rust" + if effective_engine != "rust" else () ) rust_trace: Final = ( - collect_trace(scenario_route, "rust", asynchronous=scenario.asynchronous) if engine != "python" else () + collect_trace(scenario_route, "rust", asynchronous=scenario.asynchronous) + if effective_engine != "python" + else () ) python_error: Final = _failure_message(python_trace) rust_error: Final = _failure_message(rust_trace) @@ -180,7 +185,7 @@ def execute_trace( rust: Final = pipeline_projection("rust", rust_events) except ValueError as error: return TraceArtifact.from_traces( - engine=engine, + engine=effective_engine, surface=surface, sdk_function=route.route, scenario=scenario.name, @@ -189,7 +194,7 @@ def execute_trace( python_error=f"harness: {error}", ) return TraceArtifact.from_traces( - engine=engine, + engine=effective_engine, surface=surface, sdk_function=route.route, scenario=scenario.name, diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/responses/case.py b/tests/rust-python-harness/strategies/trace_parity/sdk/responses/case.py index 68d2eaba9bf..3f6e540efc5 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/responses/case.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/responses/case.py @@ -161,7 +161,7 @@ def _anthropic_bridge_stream_fixture(engine: Engine, _base_url: str) -> RouteFix ) -SPEC: Final = RouteSpec("responses", ("responses", "aresponses"), ("responses", "aresponses"), _openai_fixture) +SPEC: Final = RouteSpec("responses", ("responses", "aresponses"), None, _openai_fixture) TRACE_SUITE: Final = TraceSuite( route=SPEC, scenarios=( diff --git a/tests/rust-python-harness/strategies/trace_parity/test_runner.py b/tests/rust-python-harness/strategies/trace_parity/test_runner.py index 57364057788..be25dd53b02 100644 --- a/tests/rust-python-harness/strategies/trace_parity/test_runner.py +++ b/tests/rust-python-harness/strategies/trace_parity/test_runner.py @@ -203,6 +203,44 @@ def test_gateway_trace_keeps_calls_outside_scenario_mappings(monkeypatch: pytest ) +def test_default_trace_skips_unavailable_rust_sdk_entrypoint(monkeypatch: pytest.MonkeyPatch) -> None: + execution: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.sdk.execution") + route: Final = RouteSpec("responses", ("responses", "aresponses"), None, _fixture) + scenario: Final = TraceScenario("sync-openai", _fixture, (), asynchronous=False) + engines: list[Engine] = [] + + def collect(_route: RouteSpec, engine: Engine, *, asynchronous: bool) -> tuple[FunctionTraceEvent, ...]: + engines.append(engine) + return (FunctionTraceEvent(0, None, "responses"),) + + monkeypatch.setattr(execution, "collect_trace", collect) + + trace: Final = execution.execute_trace(route, scenario, "sdk") + + assert engines == ["python"] + assert trace.engine == "python" + assert trace.rust_error is None + + +def test_default_trace_skips_unavailable_rust_gateway_route(monkeypatch: pytest.MonkeyPatch) -> None: + execution: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.gateway.execution") + route: Final = GatewayRouteSpec("responses", rust_supported=False) + scenario: Final = TraceScenario("async-openai", _fixture, (), asynchronous=True) + engines: list[Engine] = [] + + def collect(_route: GatewayRouteSpec, _scenario: TraceScenario, engine: Engine) -> tuple[FunctionTraceEvent, ...]: + engines.append(engine) + return (FunctionTraceEvent(0, None, "responses"),) + + monkeypatch.setattr(execution, "_collect", collect) + + trace: Final = execution.execute_gateway_trace(route, scenario) + + assert engines == ["python"] + assert trace.engine == "python" + assert trace.rust_error is None + + def test_scenario_validation_rejects_duplicate_and_unsafe_names() -> None: route: Final = RouteSpec("ocr", ("ocr", "aocr"), ("ocr", "aocr"), _fixture) duplicate: Final = TraceSuite( From db8dfe93a5a74fad5c4ba3f9fabaa5de39ad09ca Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 21:10:43 +0000 Subject: [PATCH 109/112] fix(auth): wait for the in-flight membership load before evicting its cache key on invalidation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/auth_checks.py | 10 +- .../proxy/auth/test_auth_checks.py | 95 ++++++------------- 2 files changed, 36 insertions(+), 69 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 22644bcd207..13605d7dc7b 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2204,8 +2204,6 @@ async def _fetch_team_membership_from_db( ) membership: Final = None if response is None else LiteLLM_TeamMembership.model_validate(response.dict()) _key: Final = team_membership_reservation_cache_key(user_id=user_id, team_id=team_id) - if _team_membership_inflight.get(_key) is not asyncio.current_task(): - return membership if membership is None: await user_api_key_cache.async_set_cache( key=_key, @@ -2218,8 +2216,6 @@ async def _fetch_team_membership_from_db( value=membership, model_type=LiteLLM_TeamMembership, ) - if _team_membership_inflight.get(_key) is not asyncio.current_task(): - await user_api_key_cache.async_delete_cache(key=_key) return membership @@ -2765,7 +2761,11 @@ async def invalidate_team_member_spend_state( publish_auth_cache_invalidation, ) - _team_membership_inflight.pop(team_membership_reservation_cache_key(user_id=user_id, team_id=team_id), None) + inflight: Final[object] = _team_membership_inflight.pop( + team_membership_reservation_cache_key(user_id=user_id, team_id=team_id), None + ) + if isinstance(inflight, asyncio.Task) and inflight is not asyncio.current_task(): + await asyncio.wait((inflight,)) if new_spend is not None: from litellm.proxy.proxy_server import SPEND_DB_FLOOR_CACHE_TTL_SECONDS, spend_counter_cache diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 96e9958a55d..5f87f2def93 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -6494,7 +6494,7 @@ async def test_get_team_membership_coalesces_parallel_db_fetches(): @pytest.mark.asyncio -async def test_get_team_membership_invalidation_mid_flight_discards_stale_load(): +async def test_get_team_membership_invalidation_waits_for_in_flight_load_then_evicts_it(): from litellm.proxy._types import LiteLLM_TeamMembership from litellm.proxy.auth.auth_checks import get_team_membership, invalidate_team_member_spend_state from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec @@ -6502,20 +6502,21 @@ async def test_get_team_membership_invalidation_mid_flight_discards_stale_load() started = asyncio.Event() release_stale = asyncio.Event() - release_fresh = asyncio.Event() - loads = iter((("budget-old", release_stale), ("budget-new", release_fresh))) + rows = iter(("budget-old", "budget-new")) async def _find_unique(*args, **kwargs): - budget_id, release = next(loads) + budget_id = next(rows) row = MagicMock() row.dict = lambda: {"user_id": "u-inv", "team_id": "t-inv", "spend": 1.0, "budget_id": budget_id} - started.set() - await release.wait() + if budget_id == "budget-old": + started.set() + await release_stale.wait() return row mock_prisma_client = MagicMock() mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(side_effect=_find_unique) cache = UserApiKeyCache() + _key = team_membership_reservation_cache_key(user_id="u-inv", team_id="t-inv") async def _load(): return await get_team_membership( @@ -6523,24 +6524,28 @@ async def test_get_team_membership_invalidation_mid_flight_discards_stale_load() ) stale = asyncio.create_task(_load()) - await started.wait() - await invalidate_team_member_spend_state(user_id="u-inv", team_id="t-inv", user_api_key_cache=cache) - started.clear() - fresh = asyncio.create_task(_load()) await asyncio.wait_for(started.wait(), timeout=2) - release_fresh.set() - fresh_result = await fresh - release_stale.set() - stale_result = await stale + invalidation = asyncio.create_task( + invalidate_team_member_spend_state(user_id="u-inv", team_id="t-inv", user_api_key_cache=cache) + ) + for _ in range(5): + await asyncio.sleep(0) + assert not invalidation.done() + release_stale.set() + await asyncio.wait_for(invalidation, timeout=2) + stale_result = await stale assert stale_result is not None and stale_result.budget_id == "budget-old" + assert await cache.async_get_cache(key=_key) is None + + fresh_result = await _load() assert fresh_result is not None and fresh_result.budget_id == "budget-new" assert mock_prisma_client.db.litellm_teammembership.find_unique.await_count == 2 - cached = CacheCodec.deserialize( - await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-inv", team_id="t-inv")), - model_type=LiteLLM_TeamMembership, - ) + cached = CacheCodec.deserialize(await cache.async_get_cache(key=_key), model_type=LiteLLM_TeamMembership) assert cached is not None and cached.budget_id == "budget-new" + again = await _load() + assert again is not None and again.budget_id == "budget-new" + assert mock_prisma_client.db.litellm_teammembership.find_unique.await_count == 2 @pytest.mark.asyncio @@ -6567,59 +6572,21 @@ async def test_get_team_membership_invalidation_during_cache_write_evicts_stale_ get_team_membership(user_id="u-w", team_id="t-w", prisma_client=mock_prisma_client, user_api_key_cache=cache) ) await asyncio.wait_for(write_started.wait(), timeout=2) - await invalidate_team_member_spend_state(user_id="u-w", team_id="t-w", user_api_key_cache=cache) + invalidation = asyncio.create_task( + invalidate_team_member_spend_state(user_id="u-w", team_id="t-w", user_api_key_cache=cache) + ) + for _ in range(5): + await asyncio.sleep(0) + assert not invalidation.done() + release_write.set() + await asyncio.wait_for(invalidation, timeout=2) stale_result = await stale assert stale_result is not None and stale_result.budget_id == "budget-old" assert await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-w", team_id="t-w")) is None -@pytest.mark.asyncio -async def test_get_team_membership_stale_write_finishing_after_fresh_load_never_serves_old_row(): - from litellm.proxy.auth.auth_checks import get_team_membership, invalidate_team_member_spend_state - - write_started = asyncio.Event() - release_stale_write = asyncio.Event() - stale_writes = iter((release_stale_write,)) - - class _SlowFirstWriteCache(UserApiKeyCache): - async def async_set_cache(self, key, value, local_only=False, **kwargs): - release = next(stale_writes, None) - if release is not None: - write_started.set() - await release.wait() - return await super().async_set_cache(key, value, local_only=local_only, **kwargs) - - rows = iter(("budget-old", "budget-new", "budget-new")) - - async def _find_unique(*args, **kwargs): - row = MagicMock() - row.dict = lambda: {"user_id": "u-sw", "team_id": "t-sw", "spend": 1.0, "budget_id": next(rows)} - return row - - mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(side_effect=_find_unique) - cache = _SlowFirstWriteCache() - - async def _load(): - return await get_team_membership( - user_id="u-sw", team_id="t-sw", prisma_client=mock_prisma_client, user_api_key_cache=cache - ) - - stale = asyncio.create_task(_load()) - await asyncio.wait_for(write_started.wait(), timeout=2) - await invalidate_team_member_spend_state(user_id="u-sw", team_id="t-sw", user_api_key_cache=cache) - fresh_result = await _load() - release_stale_write.set() - stale_result = await stale - after_result = await _load() - - assert fresh_result is not None and fresh_result.budget_id == "budget-new" - assert stale_result is not None and stale_result.budget_id == "budget-old" - assert after_result is not None and after_result.budget_id == "budget-new" - - @pytest.mark.asyncio async def test_common_checks_calls_get_team_membership_once_per_request(): from fastapi import Request From d2342f06ceb52f429d5a405b68adb090d76b012c Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 21:13:29 +0000 Subject: [PATCH 110/112] fix(bedrock): stamp the realtime success ownership marker when Nova Sonic spend is logged Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/bedrock/realtime/handler.py | 6 ++++- .../realtime/test_bedrock_realtime_handler.py | 27 +++++++++++++++++++ 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 7ab5a14dfc2..e83cb743aa0 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -24,7 +24,10 @@ from litellm.constants import ( from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.litellm_core_utils.realtime_streaming import DefaultLoggedRealTimeEventTypes +from litellm.litellm_core_utils.realtime_streaming import ( + REALTIME_SESSION_SUCCESS_LOGGED_KEY, + DefaultLoggedRealTimeEventTypes, +) from litellm.types.llms.openai import OpenAIRealtimeEvents from litellm.types.realtime import RealtimeResponseTransformInput @@ -346,6 +349,7 @@ class BedrockRealtime(BaseAWSLLM): prefer_async_handlers=True, ) ) + logging_obj.model_call_details[REALTIME_SESSION_SUCCESS_LOGGED_KEY] = True if outcome.provider_failure is None: return diff --git a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py index 2eadaee9e1a..7e418817e74 100644 --- a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py +++ b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py @@ -8,6 +8,7 @@ from unittest.mock import MagicMock import pytest import litellm +from litellm.litellm_core_utils.realtime_streaming import REALTIME_SESSION_SUCCESS_LOGGED_KEY from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.realtime.handler import BedrockRealtime from litellm.llms.bedrock.realtime.transformation import BedrockRealtimeConfig @@ -77,6 +78,7 @@ class UnavailableBedrockStream: class FakeLogging: def __init__(self, trace_id="trace-nova-sonic"): self.litellm_trace_id = trace_id + self.model_call_details = {} class DisconnectingClientWS: @@ -673,6 +675,31 @@ class TestBedrockRealtimeProviderFailurePropagation: await spend_dispatch["coro"] assert [event["type"] for event in spend_dispatch["events"]] == ["response.done"] + @pytest.mark.asyncio + async def test_success_dispatch_stamps_the_ownership_marker_only_when_spend_was_logged( + self, stub_aws_sdk_client, spend_dispatch + ): + stub_aws_sdk_client["streams"] = [ScriptedBedrockStream(self.TEXT_TURN)] + await BedrockRealtime().async_realtime( + model="amazon.nova-sonic-v1:0", + websocket=ConnectedClientWS([self.SESSION_UPDATE]), + logging_obj=spend_dispatch["logging_obj"], + **self.AWS_PARAMS, + ) + await spend_dispatch["coro"] + assert [event["type"] for event in spend_dispatch["events"]] == ["response.done"] + assert spend_dispatch["logging_obj"].model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY) is True + + idle_logging = FakeLogging() + stub_aws_sdk_client["streams"] = [ScriptedBedrockStream([])] + await BedrockRealtime().async_realtime( + model="amazon.nova-sonic-v1:0", + websocket=ConnectedClientWS([self.SESSION_UPDATE]), + logging_obj=idle_logging, + **self.AWS_PARAMS, + ) + assert REALTIME_SESSION_SUCCESS_LOGGED_KEY not in idle_logging.model_call_details + @pytest.mark.asyncio async def test_stream_failure_after_client_disconnect_is_not_a_provider_failure(self, stub_aws_sdk_client): stream = ScriptedBedrockStream([], receiver_type=BreakingBedrockReceiver) From b4d0f4ad2658fa2e3740d30acbd6c0ad687a3b60 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 21:33:48 +0000 Subject: [PATCH 111/112] refactor(realtime): move session ownership marker keys into constants Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 2 ++ litellm/litellm_core_utils/realtime_streaming.py | 5 +---- litellm/llms/bedrock/realtime/handler.py | 6 ++---- litellm/proxy/proxy_server.py | 7 ++----- .../litellm_core_utils/test_realtime_streaming.py | 3 +-- .../llms/bedrock/realtime/test_bedrock_realtime_handler.py | 2 +- tests/test_litellm/proxy/test_proxy_server.py | 5 +---- 7 files changed, 10 insertions(+), 20 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index e8aafda1797..c106be688e4 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -317,6 +317,8 @@ WEBSOCKET_CLOSE_REASON_MAX_BYTES: Final = 123 BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY: Final = "litellm.bedrock_realtime.pending_session_update" BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY: Final = "litellm.bedrock_realtime.session_committed" BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY: Final = "litellm.bedrock_realtime.committed_failure" +REALTIME_SESSION_SUCCESS_LOGGED_KEY: Final = "realtime_session_success_logged" +REALTIME_SESSION_FAILURE_LOGGED_KEY: Final = "realtime_session_failure_logged" # SSL/TLS cipher configuration for faster handshakes # Strategy: Strongly prefer fast modern ciphers, but allow fallback to commonly supported ones diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 2177c999804..fa567bdf4c9 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -10,6 +10,7 @@ from typing_extensions import ReadOnly import litellm from litellm._logging import redact_internal_details_from_client_message, verbose_logger +from litellm.constants import REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig from litellm.types.llms.openai import ( @@ -35,10 +36,6 @@ else: CLIENT_CONNECTION_CLASS = Any -REALTIME_SESSION_SUCCESS_LOGGED_KEY: Final = "realtime_session_success_logged" -REALTIME_SESSION_FAILURE_LOGGED_KEY: Final = "realtime_session_failure_logged" - - @dataclass(frozen=True, slots=True) class BackendClose: code: int diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index e83cb743aa0..43138b2c526 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -20,14 +20,12 @@ from litellm.constants import ( BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY, BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY, BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY, + REALTIME_SESSION_SUCCESS_LOGGED_KEY, ) from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.litellm_core_utils.realtime_streaming import ( - REALTIME_SESSION_SUCCESS_LOGGED_KEY, - DefaultLoggedRealTimeEventTypes, -) +from litellm.litellm_core_utils.realtime_streaming import DefaultLoggedRealTimeEventTypes from litellm.types.llms.openai import OpenAIRealtimeEvents from litellm.types.realtime import RealtimeResponseTransformInput diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f499ee04506..1c931863a2f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -274,6 +274,8 @@ from litellm.constants import ( PROXY_BUDGET_RESCHEDULER_MAX_TIME, PROXY_BUDGET_RESCHEDULER_MIN_TIME, PROXY_CONFIG_RELOAD_INTERVAL_SECONDS, + REALTIME_SESSION_FAILURE_LOGGED_KEY, + REALTIME_SESSION_SUCCESS_LOGGED_KEY, ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG, USER_SPEND_ALERTS_JOB_ID, WEEKLY_SPEND_REPORT_JOB_ID, @@ -12061,11 +12063,6 @@ async def realtime_websocket_endpoint( except Exception: # noqa: BLE001 # the lower layer may have closed the socket already; closing twice is not an error verbose_proxy_logger.debug("Could not close realtime client websocket; it is already gone") finally: - from litellm.litellm_core_utils.realtime_streaming import ( - REALTIME_SESSION_FAILURE_LOGGED_KEY, - REALTIME_SESSION_SUCCESS_LOGGED_KEY, - ) - if not litellm_logging_obj.model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY): await _release_realtime_budget_reservation(user_api_key_dict) if not litellm_logging_obj.model_call_details.get(REALTIME_SESSION_FAILURE_LOGGED_KEY): diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py index e1eb61b59b9..7e6d4d24905 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -10,10 +10,9 @@ from websockets.exceptions import ConnectionClosed from websockets.frames import Close import litellm +from litellm.constants import REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.realtime_streaming import ( - REALTIME_SESSION_FAILURE_LOGGED_KEY, - REALTIME_SESSION_SUCCESS_LOGGED_KEY, RealTimeStreaming, client_sent_openai_beta_realtime_header, ) diff --git a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py index 7e418817e74..6c7659d84cd 100644 --- a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py +++ b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py @@ -8,7 +8,7 @@ from unittest.mock import MagicMock import pytest import litellm -from litellm.litellm_core_utils.realtime_streaming import REALTIME_SESSION_SUCCESS_LOGGED_KEY +from litellm.constants import REALTIME_SESSION_SUCCESS_LOGGED_KEY from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.realtime.handler import BedrockRealtime from litellm.llms.bedrock.realtime.transformation import BedrockRealtimeConfig diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 64828ae3cb4..e09dddfec5b 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -10020,10 +10020,7 @@ async def _lit6973_drive_realtime_session( logging object carries a real model_call_details dict so the stamp is observable, and the reservation has empty entries so the real release touches no counter store.""" - from litellm.litellm_core_utils.realtime_streaming import ( - REALTIME_SESSION_FAILURE_LOGGED_KEY, - REALTIME_SESSION_SUCCESS_LOGGED_KEY, - ) + from litellm.constants import REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY from litellm.proxy import proxy_server as ps user_api_key_dict: Final = UserAPIKeyAuth(api_key="sk-test", token="hashed-token") From fbc1011d277b4630b4b20b061b82f5e0c0032130 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 21:42:42 +0000 Subject: [PATCH 112/112] fix(bedrock): end the realtime session when the client disconnects instead of waiting for Nova Sonic Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/bedrock/realtime/handler.py | 2 +- .../realtime/test_bedrock_realtime_handler.py | 19 +++++++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 43138b2c526..2c1ce6068b2 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -384,7 +384,7 @@ class BedrockRealtime(BaseAWSLLM): ) bedrock_task: Final = asyncio.create_task(collect_logged_events()) - await asyncio.wait((client_task, bedrock_task), return_when=asyncio.FIRST_EXCEPTION) + await asyncio.wait((client_task, bedrock_task), return_when=asyncio.FIRST_COMPLETED) client_disconnected: Final = ( client_task.done() and not client_task.cancelled() and client_task.exception() is None ) diff --git a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py index 6c7659d84cd..ac3a43b742f 100644 --- a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py +++ b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py @@ -711,6 +711,25 @@ class TestBedrockRealtimeProviderFailurePropagation: assert stream.input_stream.closed + @pytest.mark.asyncio + async def test_client_disconnect_ends_the_session_while_bedrock_output_stays_open(self, stub_aws_sdk_client): + receiver = DrainedThenOpenBedrockReceiver([]) + stream = ScriptedBedrockStream([], receiver_type=lambda _payloads: receiver) + stub_aws_sdk_client["streams"] = [stream] + + await asyncio.wait_for( + BedrockRealtime().async_realtime( + model="amazon.nova-sonic-v1:0", + websocket=RealtimeClientWS(), + logging_obj=FakeLogging(), + **self.AWS_PARAMS, + ), + timeout=1, + ) + + assert receiver.drained.is_set(), "the handler must have been waiting on the open provider stream" + assert stream.input_stream.closed + @pytest.mark.asyncio async def test_session_updated_is_not_sent_before_bedrock_is_ready(self, stub_aws_models): handler = BedrockRealtime()