[Router] Standardize route_request signature across all routers (rebase of #1023)

src/tests/test_request_auth_headers.py

@@ -32,7 +32,7 @@ def _build_request(headers=None):
     state.callbacks = None
     state.request_stats_monitor.get_request_stats.return_value = {}
     state.engine_stats_scraper.get_engine_stats.return_value = {}
-    state.router.route_request.return_value = "http://whisper-engine"
+    state.router.route_request = AsyncMock(return_value="http://whisper-engine")
 
     request = MagicMock()
     request.headers = headers or {}

src/tests/test_request_stats.py

@@ -204,8 +204,11 @@ async def test_multipart_request_counts_overlapping_requests_sharing_an_id(
     async def post(*args, **kwargs):
         return _StreamingBackendResponse()
 
+    async def route_request(*args):
+        return URL
+
     state = types.SimpleNamespace(
-        router=types.SimpleNamespace(route_request=lambda *args: URL),
+        router=types.SimpleNamespace(route_request=route_request),
         engine_stats_scraper=types.SimpleNamespace(get_engine_stats=lambda: {}),
         request_stats_monitor=monitor,
         aiohttp_client_wrapper=lambda: types.SimpleNamespace(post=post),

src/tests/test_roundrobin_router.py

@@ -62,31 +62,33 @@ def assert_even_distribution(route_counts: Counter):
     assert max(counts) - min(counts) <= 1
 
 
-def test_roundrobin_logic(
+async def test_roundrobin_logic(
     dynamic_discoveries: int = 10, max_endpoints: int = 1000, max_requests: int = 10000
 ):
     """
     Ensure that all active urls have roughly same number of requests (difference at most 1)
     """
     router = RoundRobinRouter()
 
-    def assert_router_stays_balanced(num_endpoints: int, num_requests: int):
+    async def assert_router_stays_balanced(num_endpoints: int, num_requests: int):
         endpoints, engine_stats, request_stats = generate_request_args(num_endpoints)
         route_counts = Counter()
         for _ in range(num_requests):
             request = generate_request()
-            url = router.route_request(endpoints, engine_stats, request_stats, request)
+            url = await router.route_request(
+                endpoints, engine_stats, request_stats, request, {}
+            )
             route_counts[url] += 1
 
         assert_even_distribution(route_counts)
 
     for _ in range(dynamic_discoveries):
         num_endpoints = random.randint(1, max_endpoints)
         num_requests = random.randint(1, max_requests)
-        assert_router_stays_balanced(num_endpoints, num_requests)
+        await assert_router_stays_balanced(num_endpoints, num_requests)
 
 
-def test_roundrobin_keeps_state_per_endpoint_set():
+async def test_roundrobin_keeps_state_per_endpoint_set():
     router = RoundRobinRouter()
     request = generate_request()
     engine_stats = {}
@@ -97,16 +99,20 @@ def test_roundrobin_keeps_state_per_endpoint_set():
     route_counts_b = Counter()
 
     for _ in range(100):
-        url_a = router.route_request(endpoints_a, engine_stats, request_stats, request)
+        url_a = await router.route_request(
+            endpoints_a, engine_stats, request_stats, request, {}
+        )
         route_counts_a[url_a] += 1
-        url_b = router.route_request(endpoints_b, engine_stats, request_stats, request)
+        url_b = await router.route_request(
+            endpoints_b, engine_stats, request_stats, request, {}
+        )
         route_counts_b[url_b] += 1
 
     assert_even_distribution(route_counts_a)
     assert_even_distribution(route_counts_b)
 
 
-def test_roundrobin_keeps_state_when_endpoint_order_changes():
+async def test_roundrobin_keeps_state_when_endpoint_order_changes():
     router = RoundRobinRouter()
     request = generate_request()
     engine_stats = {}
@@ -115,10 +121,11 @@ def test_roundrobin_keeps_state_when_endpoint_order_changes():
     endpoints_a_reordered = [EndpointInfo(url="a1"), EndpointInfo(url="a0")]
 
     assert (
-        router.route_request(endpoints_a, engine_stats, request_stats, request) == "a0"
+        await router.route_request(endpoints_a, engine_stats, request_stats, request)
+        == "a0"
     )
     assert (
-        router.route_request(
+        await router.route_request(
             endpoints_a_reordered,
             engine_stats,
             request_stats,
@@ -128,8 +135,8 @@ def test_roundrobin_keeps_state_when_endpoint_order_changes():
     )
 
 
-def test_roundrobin_rejects_empty_endpoint_list():
+async def test_roundrobin_rejects_empty_endpoint_list():
     router = RoundRobinRouter()
 
     with pytest.raises(ValueError, match="at least one endpoint"):
-        router.route_request([], {}, {}, generate_request())
+        await router.route_request([], {}, {}, generate_request())

src/tests/test_route_request_signature.py

@@ -0,0 +1,256 @@
+"""Tests for the standardized ``route_request`` signature (issue #1022).
+
+Every router must expose ``async def route_request(self, endpoints,
+engine_stats, request_stats, request, request_json)`` so that call sites can
+always ``await`` and always pass ``request_json`` without isinstance-based
+dispatch guards.
+"""
+
+import inspect
+import json
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import aiohttp
+import pytest
+
+from vllm_router.routers.routing_logic import (
+    DisaggregatedPrefillOrchestratedRouter,
+    DisaggregatedPrefillRouter,
+    KvawareRouter,
+    LoadAwareRouter,
+    PrefixAwareRouter,
+    PriorityRouter,
+    RoundRobinRouter,
+    RoutingInterface,
+    SessionRouter,
+)
+from vllm_router.utils import SingletonABCMeta
+
+ALL_ROUTER_CLASSES = [
+    RoundRobinRouter,
+    SessionRouter,
+    KvawareRouter,
+    LoadAwareRouter,
+    PrefixAwareRouter,
+    PriorityRouter,
+    DisaggregatedPrefillRouter,
+    DisaggregatedPrefillOrchestratedRouter,
+]
+
+
+class EndpointInfo:
+    def __init__(self, url, model_names=None, sleep=False, Id=None):
+        self.url = url
+        self.model_names = model_names or ["test-model"]
+        self.sleep = sleep
+        self.Id = Id
+
+
+class RecordingAsyncRouter(RoutingInterface):
+    """A router that is NOT in any isinstance whitelist at the call sites.
+
+    It records the ``request_json`` it receives, so tests can verify that the
+    call sites actually awaited the coroutine (the body of an un-awaited
+    coroutine never runs) and passed the request body through.
+    """
+
+    def __init__(self):
+        if hasattr(self, "_initialized"):
+            return
+        self.received_request_jsons = []
+        self.max_instance_failover_reroute_attempts = 0
+        self._initialized = True
+
+    async def route_request(
+        self,
+        endpoints,
+        engine_stats,
+        request_stats,
+        request,
+        request_json=None,
+    ) -> str:
+        self.received_request_jsons.append(request_json)
+        return endpoints[0].url
+
+
+@pytest.fixture(autouse=True)
+def cleanup_singletons():
+    yield
+    for cls in list(SingletonABCMeta._instances.keys()):
+        del SingletonABCMeta._instances[cls]
+
+
+@pytest.mark.parametrize("router_cls", ALL_ROUTER_CLASSES)
+def test_route_request_is_coroutine_function(router_cls):
+    assert inspect.iscoroutinefunction(
+        router_cls.route_request
+    ), f"{router_cls.__name__}.route_request must be an async def"
+
+
+@pytest.mark.parametrize("router_cls", ALL_ROUTER_CLASSES)
+def test_route_request_accepts_request_json(router_cls):
+    params = inspect.signature(router_cls.route_request).parameters
+    assert "request_json" in params, (
+        f"{router_cls.__name__}.route_request must accept a request_json" " parameter"
+    )
+    param_names = list(params)
+    assert param_names[:6] == [
+        "self",
+        "endpoints",
+        "engine_stats",
+        "request_stats",
+        "request",
+        "request_json",
+    ], f"{router_cls.__name__}.route_request has unexpected parameter order"
+    assert (
+        params["request_json"].default is None
+    ), f"{router_cls.__name__}.route_request must default request_json to None"
+
+
+def test_base_interface_route_request_is_async_with_request_json():
+    assert inspect.iscoroutinefunction(RoutingInterface.route_request)
+    assert (
+        "request_json" in inspect.signature(RoutingInterface.route_request).parameters
+    )
+
+
+@pytest.mark.asyncio
+async def test_roundrobin_routes_via_awaited_path():
+    router = RoundRobinRouter()
+    endpoints = [EndpointInfo(url="http://engine1"), EndpointInfo(url="http://engine2")]
+    request = MagicMock()
+
+    urls = [
+        await router.route_request(endpoints, {}, {}, request, {}) for _ in range(4)
+    ]
+
+    assert urls == ["http://engine1", "http://engine2"] * 2
+
+
+@pytest.mark.asyncio
+async def test_route_general_request_awaits_any_router_and_passes_request_json():
+    """Call-site dispatch must not depend on the router's concrete type."""
+    router = RecordingAsyncRouter()
+
+    sd = MagicMock()
+    sd.get_endpoint_info.return_value = [
+        EndpointInfo(url="http://engine1"),
+        EndpointInfo(url="http://engine2"),
+    ]
+    sd.aliases = None
+    sd.has_ever_seen_model.return_value = True
+
+    state = MagicMock()
+    state.router = router
+    state.engine_stats_scraper.get_engine_stats.return_value = {}
+    state.request_stats_monitor.get_request_stats.return_value = {}
+    state.otel_enabled = False
+    state.semantic_cache_available = False
+    state.callbacks = None
+    state.external_provider_registry = None
+
+    req = MagicMock()
+    req.headers = {"content-type": "application/json"}
+    req.query_params = {}
+    req.method = "POST"
+    req.url = "http://router/v1/chat/completions"
+    req.app.state = state
+
+    request_body = {"model": "test-model", "stream": False}
+
+    async def body():
+        return json.dumps(request_body).encode()
+
+    req.body = body
+
+    mock_headers = MagicMock()
+    mock_headers.items.return_value = [("content-type", "text/event-stream")]
+
+    async def ok(*args, **kwargs):
+        yield mock_headers, 200
+        yield b"done"
+
+    with (
+        patch(
+            "vllm_router.services.request_service.request.get_service_discovery",
+            return_value=sd,
+        ),
+        patch(
+            "vllm_router.services.request_service.request.is_request_rewriter_initialized",
+            return_value=False,
+        ),
+        patch(
+            "vllm_router.services.request_service.request.process_request",
+            side_effect=ok,
+        ) as mock_process,
+    ):
+        from vllm_router.services.request_service.request import route_general_request
+
+        resp = await route_general_request(req, "/v1/chat/completions", MagicMock())
+
+    assert resp.status_code == 200
+    # The coroutine body only runs when awaited, so this proves the call site
+    # awaited route_request and forwarded the parsed JSON body.
+    assert router.received_request_jsons == [request_body]
+    # process_request must receive the resolved URL string, not a coroutine.
+    assert mock_process.call_args.args[2] == "http://engine1"
+
+
+@pytest.mark.asyncio
+async def test_proxy_multipart_request_awaits_route_request_with_request_json():
+    router = RecordingAsyncRouter()
+
+    state = MagicMock()
+    state.router = router
+    state.engine_stats_scraper.get_engine_stats.return_value = {}
+    state.request_stats_monitor.get_request_stats.return_value = {}
+    state.otel_enabled = False
+
+    req = MagicMock()
+    req.headers = {"content-type": "multipart/form-data; boundary=router-boundary"}
+    req.query_params = {}
+    req.method = "POST"
+    req.url = "http://router/v1/audio/transcriptions"
+    req.app.state = state
+
+    backend_response = MagicMock()
+    backend_response.status = 200
+    backend_response.headers.items.return_value = [("content-type", "application/json")]
+    backend_response.json = AsyncMock(return_value={"ok": True})
+
+    mock_client = MagicMock()
+    mock_client.post = AsyncMock(return_value=backend_response)
+    req.app.state.aiohttp_client_wrapper = MagicMock(return_value=mock_client)
+
+    sd = MagicMock()
+    sd.get_endpoint_info.return_value = [
+        EndpointInfo(url="http://engine1", model_names=["whisper-model"])
+    ]
+
+    form_data = aiohttp.FormData()
+    form_data.add_field("file", b"audio-bytes", filename="sample.wav")
+    form_data.add_field("model", "whisper-model")
+
+    with patch(
+        "vllm_router.services.request_service.request.get_service_discovery",
+        return_value=sd,
+    ):
+        from vllm_router.services.request_service.request import (
+            proxy_multipart_request,
+        )
+
+        response = await proxy_multipart_request(
+            form_data,
+            "whisper-model",
+            "/v1/audio/transcriptions",
+            req,
+        )
+
+    assert response.status_code == 200
+    # Proves the call site awaited route_request and passed a request_json
+    # ({} because multipart requests carry no JSON body).
+    assert router.received_request_jsons == [{}]
+    # And the backend was called with the resolved URL string.
+    assert (
+        mock_client.post.await_args.args[0] == "http://engine1/v1/audio/transcriptions"
+    )