From a55ff8d83ed4e4a82f0cd95c67c09a691e1244d7 Mon Sep 17 00:00:00 2001
From: tin <tin@berri.ai>
Date: Thu, 6 Aug 2026 04:15:47 +0000
Subject: [PATCH] fix(router): finalize complexity affinity cleanup

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 litellm/router.py                             | 37 +++++++------------
 .../deployment_affinity_check.py              |  2 +-
 .../test_session_id_affinity.py               | 23 ++++++++++++
 3 files changed, 37 insertions(+), 25 deletions(-)

diff --git a/litellm/router.py b/litellm/router.py
index 81c80573c66..249039fb8ba 100644
--- a/litellm/router.py
+++ b/litellm/router.py
@@ -21,7 +21,7 @@ import traceback
 import weakref
 from collections import defaultdict
 from collections.abc import AsyncGenerator, Callable, Generator, Iterator, Mapping, Sequence
-from functools import lru_cache, reduce
+from functools import lru_cache
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeVar, Union, cast
 
@@ -269,20 +269,6 @@ def _iter_complexity_router_session_affinity_groups(
         yield config.default_model, config.session_affinity_ttl_seconds
 
 
-def _merge_minimum_session_affinity_ttl(
-    group_ttls: Mapping[str, int],
-    group_ttl: tuple[str, int],
-) -> Mapping[str, int]:
-    model_group, ttl = group_ttl
-    existing_ttl: Final[int | None] = group_ttls.get(model_group)
-    return MappingProxyType(
-        {
-            **group_ttls,
-            model_group: ttl if existing_ttl is None else min(existing_ttl, ttl),
-        }
-    )
-
-
 def _cost_value_as_float(value: str | float | None) -> float | None:
     if value is None:
         return None
@@ -7651,15 +7637,18 @@ class Router:
         return classify_strategy_router_model(litellm_params.model) == "complexity"
 
     def _get_complexity_router_session_affinity_group_ttls(self) -> Mapping[str, int]:
-        return reduce(
-            _merge_minimum_session_affinity_ttl,
-            (
-                group_ttl
-                for strategies in self.complexity_routers.values()
-                for tagged_strategy in strategies
-                for group_ttl in _iter_complexity_router_session_affinity_groups(tagged_strategy.strategy)
-            ),
-            MappingProxyType({}),
+        entries: Final = tuple(
+            group_ttl
+            for strategies in self.complexity_routers.values()
+            for tagged_strategy in strategies
+            for group_ttl in _iter_complexity_router_session_affinity_groups(tagged_strategy.strategy)
+        )
+        groups: Final = frozenset(model_group for model_group, _ in entries)
+        return MappingProxyType(
+            {
+                model_group: min(ttl for candidate_group, ttl in entries if candidate_group == model_group)
+                for model_group in groups
+            }
         )
 
     def _ensure_deployment_affinity_check(self) -> None:
diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py
index 30071ae4a57..6a1a4156ee3 100644
--- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py
+++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py
@@ -463,7 +463,7 @@ class DeploymentAffinityCheck(CustomLogger):
             self._get_session_id_from_request_kwargs(request_kwargs=kwargs) if enable_session_id else None
         )
 
-        if (enable_user_key and user_key is None) and (enable_session_id and session_id is None):
+        if not ((enable_user_key and user_key is not None) or (enable_session_id and session_id is not None)):
             return None
 
         model_info = kwargs.get("model_info")
diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py b/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py
index 3ad0a79af67..37521701434 100644
--- a/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py
+++ b/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py
@@ -347,6 +347,29 @@ async def test_complexity_router_session_affinity_uses_router_configured_ttl():
     assert any(call.kwargs.get("ttl") == 17 for call in cache.async_set_cache.call_args_list)
 
 
+@pytest.mark.asyncio
+async def test_session_affinity_without_session_id_does_not_write_cache():
+    router = _complexity_router()
+    callback = next(
+        callback for callback in router.optional_callbacks or [] if isinstance(callback, DeploymentAffinityCheck)
+    )
+    cache = AsyncMock()
+    callback.cache = cache
+
+    await callback.async_pre_call_deployment_hook(
+        kwargs={
+            "model_info": {"id": "deployment-1"},
+            "metadata": {
+                "deployment_model_name": "target-group",
+                "user_api_key_hash": "key-1",
+            },
+        },
+        call_type=None,
+    )
+
+    cache.async_set_cache.assert_not_called()
+
+
 @pytest.mark.asyncio
 async def test_complexity_router_session_affinity_expires_and_reselects():
     router = _complexity_router(session_affinity_ttl_seconds=1)