fix: logic extend on total_load AND loaded_count

test/test_choose_endpoint.py

@@ -112,6 +112,37 @@ class TestChooseEndpointBasic:
                 assert ep in (EP1, EP2)
                 assert ep != EP3
 
+    async def test_two_cold_models_spread_across_backends(self):
+        # Regression: 3 backends all advertise all models. Two *different*
+        # cold models requested back-to-back must land on *different*
+        # backends. Once model-a is resident on the chosen backend (infinite
+        # keep-alive), its in-flight count drops back to 0 — so only the
+        # resident-model count distinguishes the backends. Without it, the
+        # second cold model would randomly re-collide on the busy backend.
+        cfg = _make_cfg([EP1, EP2, EP3], max_conn=4)
+
+        async def available(ep, *_):
+            return {"model-a:latest", "model-b:latest"}
+
+        # model-a finished loading on EP1 and stays resident; its request has
+        # completed so EP1 has zero in-flight load, same as EP2/EP3.
+        loaded = {EP1: {"model-a:latest"}, EP2: set(), EP3: set()}
+
+        async def loaded_models(ep):
+            return loaded[ep]
+
+        with (
+            patch.object(router, "config", cfg),
+            patch.object(router.fetch, "available_models", side_effect=available),
+            patch.object(router.fetch, "loaded_models", side_effect=loaded_models),
+        ):
+            # A cold model-b must avoid EP1 (which already holds model-a) and
+            # go to one of the empty backends, every time.
+            for _ in range(50):
+                ep, _ = await router.choose_endpoint("model-b:latest", reserve=False)
+                assert ep in (EP2, EP3)
+                assert ep != EP1
+
     async def test_saturated_picks_least_busy(self):
         cfg = _make_cfg([EP1, EP2])
         cfg.max_concurrent_connections = 1