diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index a09b6701e8d..4ae17785804 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -79,6 +79,7 @@ class _ClientWebSocket(_ScopedWebSocket, Protocol): _SESSION_FIELDS: Final = TypeAdapter(dict[str, object]) _NO_SESSION_FIELDS: Final[Mapping[str, object]] = MappingProxyType({}) +_SESSION_MERGE_MAX_DEPTH: Final = 8 class _SessionUpdateFrame(TypedDict): @@ -86,15 +87,17 @@ class _SessionUpdateFrame(TypedDict): session: ReadOnly[Mapping[str, object]] -def _merge_session_fields(base: Mapping[str, object], update: Mapping[str, object]) -> Mapping[str, object]: +def _merge_session_fields( + base: Mapping[str, object], update: Mapping[str, object], depth: int = 0 +) -> Mapping[str, object]: def merged_value(key: str) -> object: if key not in update: return base[key] value: Final = update[key] base_value: Final = base.get(key) - if isinstance(value, Mapping) and isinstance(base_value, Mapping): + if depth < _SESSION_MERGE_MAX_DEPTH and isinstance(value, Mapping) and isinstance(base_value, Mapping): return _merge_session_fields( - _SESSION_FIELDS.validate_python(base_value), _SESSION_FIELDS.validate_python(value) + _SESSION_FIELDS.validate_python(base_value), _SESSION_FIELDS.validate_python(value), depth + 1 ) return value diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index e9f87ba6cae..4ac150d788b 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -67,6 +67,7 @@ IGNORE_FUNCTIONS = [ "_redact_agent_params_tree", # max depth set (default 10), same shape as _redact_sensitive_litellm_params. "_restore_redacted_nested_value", # max depth set (default 10), mirrors _redact_agent_params_tree on the write side. "_unqualified", # bounded by the qualifier depth of a static TypedDict annotation (Annotated, Required/NotRequired, ReadOnly around one type, no cycles possible). + "_merge_session_fields", # max depth set (_SESSION_MERGE_MAX_DEPTH); deeper mappings are replaced wholesale by the later session.update. ] 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 859afa01e96..a4ed017d7df 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -10,6 +10,7 @@ import litellm from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.realtime_streaming import ( RealTimeStreaming, + _merge_session_fields, client_sent_openai_beta_realtime_header, ) from litellm.llms.xai.realtime.transformation import XAIRealtimeNormalizer @@ -2462,6 +2463,20 @@ async def test_default_mode_merges_nested_audio_settings_across_session_updates( assert setup["generationConfig"]["speechConfig"]["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] == "Kore" +def test_merge_session_fields_replaces_nested_mappings_wholesale_past_the_depth_cap(): + def nested(levels: int, key: str) -> dict: + return {key: levels, **({"child": nested(levels - 1, key)} if levels else {})} + + def level(node, depth: int): + return node if depth == 0 else level(node["child"], depth - 1) + + merged = _merge_session_fields(nested(10, "keep"), nested(10, "new")) + + assert all({"keep", "new"} <= set(level(merged, depth)) for depth in range(9)) + assert "keep" not in level(merged, 9) + assert level(merged, 9)["new"] == 1 + + @pytest.mark.asyncio async def test_default_mode_acknowledges_each_held_session_update_before_any_backend_frame(monkeypatch): """A client that waits for session.updated before sending content must not