From f80cb5cb4676a9d49009aef432cc9f5eda5ed15b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 15 Sep 2026 00:04:02 -0700 Subject: [PATCH] fix(router): ignore planted request_retry_count seeds and cover the rust OCR cap path The router clamps a negative request_retry_count found in request metadata before counting a failure, and the proxy strips a client-supplied request_retry_count with the other router-reserved metadata fields. The rust OCR lifecycle test that trips the per-request cap now plants request_retry_count instead of attempted_retries, which the cap no longer reads since the previous commit --- litellm/proxy/litellm_pre_call_utils.py | 2 +- litellm/router.py | 4 ++-- .../test_router_helper_utils.py | 6 ++++-- .../proxy/test_litellm_pre_call_utils.py | 4 ++++ tests/test_litellm/test_router.py | 16 ++++++++++++---- tests/test_litellm_rust/ocr/test_lifecycle.py | 2 +- 6 files changed, 24 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 59971e54e46..14d0e7e478f 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -334,7 +334,7 @@ _CLIENT_PRICING_METADATA_FIELDS: Final = frozenset({"model_info", "standard_logg # and read by spend logs as fact; a client value has no legitimate meaning and no # key or team setting keeps it, so the strip is never gated. _ROUTER_RESERVED_METADATA_FIELDS: Final = frozenset( - {"attempted_fallbacks", "original_model_group", CLIENT_OUTPUT_CEILING_METADATA_KEY} + {"attempted_fallbacks", "original_model_group", "request_retry_count", CLIENT_OUTPUT_CEILING_METADATA_KEY} ) _ALLOW_CLIENT_PRICING_OVERRIDE_METADATA_KEY: Final = "allow_client_pricing_override" diff --git a/litellm/router.py b/litellm/router.py index ff50cac1328..cb0b7050876 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -8403,8 +8403,8 @@ class Router: else () ) breadcrumbs: Final = (*kept_breadcrumbs, attempt_record) - earlier_retry_count: Final = request_metadata.get("request_retry_count") - request_retry_count: Final = (earlier_retry_count if type(earlier_retry_count) is int else 0) + 1 + earlier: Final = request_metadata.get("request_retry_count") + request_retry_count: Final = (earlier if type(earlier) is int and 0 <= earlier else 0) + 1 kwargs[_metadata_var]["previous_models"] = breadcrumbs # rebind-ok: the logging object already holds this dict kwargs[_metadata_var]["request_retry_count"] = request_retry_count # rebind-ok: same dict, read by the cap return kwargs diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index c7e577366c3..7949dd7818d 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -632,8 +632,8 @@ def test_deployment_callback_respects_cooldown_time(model_list): @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, copies neither the request kwargs nor the - request metadata into it, and counts every failed attempt of the request independently of the - per-hop attempted_retries""" + request metadata into it, counts every failed attempt of the request independently of the + per-hop attempted_retries, and never trusts a negative count planted before the first failure""" router = Router(model_list=model_list) rate_limit_error = litellm.RateLimitError(message="slow down", llm_provider="openai", model="gpt-3.5-turbo") new_kwargs = router.log_retry( @@ -656,6 +656,8 @@ def test_log_retry(model_list, metadata_key): ] assert new_kwargs[metadata_key]["request_retry_count"] == 1 assert router.log_retry(kwargs=new_kwargs, e=rate_limit_error)[metadata_key]["request_retry_count"] == 2 + planted_kwargs = {"model": "gpt-3.5-turbo", metadata_key: {"request_retry_count": -100}} + assert router.log_retry(kwargs=planted_kwargs, e=rate_limit_error)[metadata_key]["request_retry_count"] == 1 def test_update_usage(model_list): diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 37b983d709a..abd59bb4d3b 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -7346,6 +7346,7 @@ def _reserved_stamp_key(key_metadata: dict | None = None) -> UserAPIKeyAuth: _PLANTED_STAMPS = { "attempted_fallbacks": 99, "original_model_group": "spoofed-group", + "request_retry_count": -100, "_client_output_ceiling": {"api_base": "https://attacker.example"}, "client_key": "client_value", } @@ -7378,6 +7379,7 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_from_bo assert "attempted_fallbacks" not in updated["metadata"] assert "original_model_group" not in updated["metadata"] assert "_client_output_ceiling" not in updated["metadata"] + assert "request_retry_count" not in updated["metadata"] assert updated["metadata"]["client_key"] == "client_value" @@ -7403,6 +7405,7 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_from_js assert "litellm_metadata" not in updated assert "attempted_fallbacks" not in updated["metadata"] assert "original_model_group" not in updated["metadata"] + assert "request_retry_count" not in updated["metadata"] assert updated["metadata"]["client_key"] == "client_value" @@ -7431,6 +7434,7 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_despite assert updated["metadata"]["model_info"] == {"input_cost_per_token": 0.0} assert "attempted_fallbacks" not in updated["metadata"] assert "original_model_group" not in updated["metadata"] + assert "request_retry_count" not in updated["metadata"] @pytest.mark.asyncio diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 577332727d7..1de86084980 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -11089,18 +11089,26 @@ def _failing_group_with_healthy_fallback_router(num_retries): @pytest.mark.asyncio @pytest.mark.parametrize( - "cap, hop_refused", [(2, True), (4, False)], ids=["cap-spent-before-the-hop", "cap-not-reached-by-the-hop"] + "cap, planted_count, hop_refused", + [(2, None, True), (4, None, False), (2, -100, True)], + ids=["cap-spent-before-the-hop", "cap-not-reached-by-the-hop", "planted-negative-count-does-not-lift-the-cap"], ) -async def test_num_retries_per_request_counts_retries_across_fallback_hops(monkeypatch, cap, hop_refused): +async def test_num_retries_per_request_counts_retries_across_fallback_hops( + monkeypatch, cap, planted_count, hop_refused +): """num_retries_per_request caps the retries of one request, fallback hops included. Each hop starts a fresh per-hop attempted_retries at zero, so a cap read from that counter let every hop retry from zero - and a request could spend far more retries than the cap allows.""" + and a request could spend far more retries than the cap allows. A caller who plants a negative count + in the request metadata must not push the cap further away either.""" monkeypatch.setattr(litellm, "num_retries_per_request", cap) router = _failing_group_with_healthy_fallback_router(num_retries=1) recorder = _FallbackAttemptRecorder() litellm.callbacks.append(recorder) try: - request = router.acompletion(model="broken-group", messages=[{"role": "user", "content": "hi"}]) + metadata = {} if planted_count is None else {"request_retry_count": planted_count} + request = router.acompletion( + model="broken-group", messages=[{"role": "user", "content": "hi"}], metadata=metadata + ) if not hop_refused: assert (await request).choices[0].message.content == "ok" return diff --git a/tests/test_litellm_rust/ocr/test_lifecycle.py b/tests/test_litellm_rust/ocr/test_lifecycle.py index 77d9ef167d0..dfcd63d3019 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": {"attempted_retries": 1}} + arguments: Final = {"document": {"type": "file", "file": File()}, "metadata": {"request_retry_count": 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 == []