From 272458be0cf99b09d8ab596acfb838a3bfe83a70 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Thu, 27 Aug 2026 18:03:25 -0700 Subject: [PATCH] fix(router): copy instead of mutating caller metadata when scrubbing fallback stamp keys (#38586) --- litellm/router.py | 10 +- .../router_utils/fallback_event_handlers.py | 14 +- tests/test_litellm/test_router.py | 223 +++++++++++++++++- 3 files changed, 233 insertions(+), 14 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 021dafa9791..0afefdbf7cd 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -6982,9 +6982,13 @@ class Router: _sibling_metadata_key: Final = ( "metadata" if _fallback_metadata_key == "litellm_metadata" else "litellm_metadata" ) - if isinstance(_sibling_metadata := kwargs.get(_sibling_metadata_key), dict): - _sibling_metadata.pop("attempted_fallbacks", None) - _sibling_metadata.pop("original_model_group", None) + if isinstance(_sibling_metadata := kwargs.get(_sibling_metadata_key), dict) and ( + "attempted_fallbacks" in _sibling_metadata or "original_model_group" in _sibling_metadata + ): + _scrubbed_sibling_metadata: Final = _sibling_metadata.copy() + _scrubbed_sibling_metadata.pop("attempted_fallbacks", None) + _scrubbed_sibling_metadata.pop("original_model_group", None) + kwargs[_sibling_metadata_key] = _scrubbed_sibling_metadata if isinstance(_fallback_metadata := kwargs.get(_fallback_metadata_key), dict): _fallback_metadata["attempted_fallbacks"] = 0 if model_group is not None: diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index acdc7df5bd1..af7efa51e3f 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -375,12 +375,14 @@ async def run_async_fallback( elif isinstance(mg, dict): kwargs.update(mg) fallback_depth = fallback_depth + 1 - kwargs[metadata_variable_name] = { - "original_model_group": original_model_group, - **(kwargs.get(metadata_variable_name) or {}), - "model_group": kwargs.get("model", None), - "attempted_fallbacks": fallback_depth, - } + _hop_metadata = dict(kwargs.get(metadata_variable_name) or {}) + _original_model_group_stamp = _hop_metadata.pop("original_model_group", original_model_group) + _hop_metadata.pop("model_group", None) + _hop_metadata.pop("attempted_fallbacks", None) + _hop_metadata["original_model_group"] = _original_model_group_stamp + _hop_metadata["model_group"] = kwargs.get("model", None) + _hop_metadata["attempted_fallbacks"] = fallback_depth + kwargs[metadata_variable_name] = _hop_metadata kwargs["fallback_depth"] = fallback_depth kwargs["max_fallbacks"] = max_fallbacks kwargs["attempted_targets"] = attempted diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 8716e6d6b25..1517ac0e6d0 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -1,5 +1,6 @@ import asyncio import copy +import functools import json import logging import os @@ -7,6 +8,7 @@ import threading from unittest.mock import AsyncMock, MagicMock, patch import httpx +import openai import pytest @@ -10595,11 +10597,26 @@ async def test_async_function_with_fallbacks_skips_stamp_on_genuine_reentrant_ho assert metadata["original_model_group"] == "prod-chat" +def _record_router_acompletion_kwargs(router: litellm.Router) -> list: + """Spy on router._acompletion, recording each call's kwargs while delegating through.""" + records = [] + original_acompletion = router._acompletion + + @functools.wraps(original_acompletion) + async def _spy(*args, **spy_kwargs): + records.append(spy_kwargs) + return await original_acompletion(*args, **spy_kwargs) + + router._acompletion = _spy + return records + + @pytest.mark.asyncio async def test_async_function_with_fallbacks_scrubs_spoofed_values_from_sibling_bucket(): """Spend logs read a truthy litellm_metadata dict in preference to metadata, so spoofed - stamp keys planted in the bucket the route does not own are removed on entry instead of - flowing into the spend log row.""" + stamp keys planted in the bucket the route does not own are removed from the request's + downstream view on entry instead of flowing into the spend log row. The caller's own + dict object is never mutated: the scrub replaces the kwargs entry with a cleaned copy.""" router = litellm.Router( model_list=[ { @@ -10614,6 +10631,7 @@ async def test_async_function_with_fallbacks_scrubs_spoofed_values_from_sibling_ "original_model_group": "spoofed-group", "client_key": "client_value", } + downstream_calls = _record_router_acompletion_kwargs(router) await router.acompletion( model="gpt-3.5-turbo", @@ -10622,13 +10640,208 @@ async def test_async_function_with_fallbacks_scrubs_spoofed_values_from_sibling_ litellm_metadata=litellm_metadata, ) - assert "attempted_fallbacks" not in litellm_metadata - assert "original_model_group" not in litellm_metadata - assert litellm_metadata["client_key"] == "client_value" + assert len(downstream_calls) == 1 + downstream_sibling = downstream_calls[0]["litellm_metadata"] + assert "attempted_fallbacks" not in downstream_sibling + assert "original_model_group" not in downstream_sibling + assert downstream_sibling["client_key"] == "client_value" + assert litellm_metadata == { + "attempted_fallbacks": 99, + "original_model_group": "spoofed-group", + "client_key": "client_value", + } assert metadata["attempted_fallbacks"] == 0 assert metadata["original_model_group"] == "gpt-3.5-turbo" +@pytest.mark.asyncio +async def test_async_function_with_fallbacks_leaves_caller_sibling_dict_object_untouched(): + """The sibling-bucket scrub hands downstream a cleaned copy and never edits the dict + object the caller passed in: callers reuse metadata dicts across requests, and logging + callbacks observe the caller's object.""" + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "hi"}, + } + ] + ) + litellm_metadata = { + "attempted_fallbacks": 7, + "original_model_group": "planted-group", + "client_key": "client_value", + } + caller_snapshot = copy.deepcopy(litellm_metadata) + downstream_calls = _record_router_acompletion_kwargs(router) + + await router.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hey"}], + metadata={}, + litellm_metadata=litellm_metadata, + ) + + assert len(downstream_calls) == 1 + assert downstream_calls[0]["litellm_metadata"] is not litellm_metadata + assert litellm_metadata == caller_snapshot + + +@pytest.mark.asyncio +async def test_async_function_with_fallbacks_passes_clean_sibling_bucket_through_unchanged(): + """A sibling bucket carrying no reserved stamp keys is forwarded downstream as the + caller's own object with no copy made, matching pre-scrub behavior. Retry accounting + stamped into that bucket downstream predates the scrub and is out of its scope.""" + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "hi"}, + } + ] + ) + litellm_metadata = {"client_key": "client_value"} + downstream_calls = _record_router_acompletion_kwargs(router) + + await router.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hey"}], + metadata={}, + litellm_metadata=litellm_metadata, + ) + + assert len(downstream_calls) == 1 + assert downstream_calls[0]["litellm_metadata"] is litellm_metadata + assert litellm_metadata["client_key"] == "client_value" + assert "attempted_fallbacks" not in litellm_metadata + assert "original_model_group" not in litellm_metadata + + +@pytest.mark.asyncio +async def test_run_async_fallback_keeps_caller_metadata_keys_on_the_wire(monkeypatch): + """Under enable_preview_features, add_openai_metadata forwards only the first 16 + string pairs of request metadata to the provider body, so the fallback hop must + spread caller keys before the router's own stamps: a stamp inserted first evicts + the caller's 16th key from the wire while the internal stamp rides in its place.""" + monkeypatch.setattr(litellm, "enable_preview_features", True) + caller_metadata = {f"user_key_{i}": f"value_{i}" for i in range(16)} + router = litellm.Router( + model_list=[ + { + "model_name": "primary-group", + "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "sk-test"}, + }, + { + "model_name": "fallback-group", + "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "sk-test"}, + }, + ], + fallbacks=[{"primary-group": ["fallback-group"]}], + num_retries=0, + ) + + wire_bodies = [] + + def _respond(request: httpx.Request) -> httpx.Response: + wire_bodies.append(json.loads(request.content)) + return httpx.Response( + 200, + json={ + "id": "chatcmpl-wire", + "object": "chat.completion", + "created": 1, + "model": "gpt-3.5-turbo", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + + client = openai.AsyncOpenAI( + api_key="sk-test", + http_client=httpx.AsyncClient(transport=httpx.MockTransport(_respond)), + ) + + await router.acompletion( + model="primary-group", + messages=[{"role": "user", "content": "hey"}], + metadata=dict(caller_metadata), + mock_testing_fallbacks=True, + client=client, + ) + + assert len(wire_bodies) == 1 + assert wire_bodies[0]["metadata"] == caller_metadata + + wire_bodies.clear() + small_metadata = {"team": "alpha", "env": "prod"} + await router.acompletion( + model="primary-group", + messages=[{"role": "user", "content": "hey again"}], + metadata=dict(small_metadata), + mock_testing_fallbacks=True, + client=client, + ) + + assert len(wire_bodies) == 1 + small_wire = wire_bodies[0]["metadata"] + assert {k: small_wire[k] for k in small_metadata} == small_metadata + assert small_wire["original_model_group"] == "primary-group" + assert small_wire["model_group"] == "fallback-group" + + +@pytest.mark.asyncio +async def test_run_async_fallback_two_hop_chain_reports_entry_group_and_hop_count(): + """A two-hop fallback chain stamps attempted_fallbacks=2 on the final leg and keeps + original_model_group at the group requested on entry: a later hop's stamp appends + after caller keys without overriding the value stamped by an earlier hop.""" + router = litellm.Router( + model_list=[ + { + "model_name": "group-a", + "litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "litellm.InternalServerError"}, + }, + { + "model_name": "group-b", + "litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "litellm.InternalServerError"}, + }, + { + "model_name": "group-c", + "litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "ok"}, + }, + ], + fallbacks=[{"group-a": ["group-b"]}, {"group-b": ["group-c"]}], + num_retries=0, + ) + metadata = {} + leg_records = [] + original_acompletion = router._acompletion + + @functools.wraps(original_acompletion) + async def _spy(*args, **spy_kwargs): + leg_records.append((spy_kwargs.get("model"), copy.deepcopy(spy_kwargs.get("metadata")))) + return await original_acompletion(*args, **spy_kwargs) + + router._acompletion = _spy + + await router.acompletion( + model="group-a", + messages=[{"role": "user", "content": "hey"}], + metadata=metadata, + ) + + assert [model for model, _ in leg_records] == ["group-a", "group-b", "group-c"] + hop_one_metadata = leg_records[1][1] + assert hop_one_metadata["attempted_fallbacks"] == 1 + assert hop_one_metadata["original_model_group"] == "group-a" + assert hop_one_metadata["model_group"] == "group-b" + hop_two_metadata = leg_records[2][1] + assert hop_two_metadata["attempted_fallbacks"] == 2 + assert hop_two_metadata["original_model_group"] == "group-a" + assert hop_two_metadata["model_group"] == "group-c" + assert metadata["attempted_fallbacks"] == 0 + assert metadata["original_model_group"] == "group-a" + + def _permission_denied_error() -> litellm.PermissionDeniedError: return litellm.PermissionDeniedError( message="OpenrouterException - this key has no access to the model",