ray-project--ray/python/ray/serve/tests/unit/test_router.py

import asyncio
import concurrent.futures
import random
import sys
import threading
from collections import defaultdict
from contextlib import asynccontextmanager
from dataclasses import replace
from typing import Callable, Dict, List, Optional, Set, Tuple
from unittest.mock import Mock, patch

import grpc
import pytest

import ray
from ray._common.test_utils import async_wait_for_condition, wait_for_condition
from ray._common.utils import get_or_create_event_loop
from ray.exceptions import (
    ActorDiedError,
    ActorUnavailableError,
    RayTaskError,
    TaskCancelledError,
)
from ray.serve._private.common import (
    DeploymentHandleSource,
    DeploymentID,
    ReplicaID,
    ReplicaQueueLengthInfo,
    RequestMetadata,
    RunningReplicaInfo,
)
from ray.serve._private.config import DeploymentConfig
from ray.serve._private.constants import (
    RAY_SERVE_COLLECT_AUTOSCALING_METRICS_ON_HANDLE,
    RAY_SERVE_METRICS_EXPORT_INTERVAL_MS,
)
from ray.serve._private.replica import Replica as ServeReplica
from ray.serve._private.replica_result import ReplicaResult, gRPCReplicaResult
from ray.serve._private.request_router import (
    PendingRequest,
    RequestRouter,
    RunningReplica,
)
from ray.serve._private.request_router.common import ReplicaQueueLengthCache
from ray.serve._private.request_router.replica_wrapper import ReplicaSelection
from ray.serve._private.router import (
    QUEUED_REQUESTS_KEY,
    AsyncioRouter,
    CurrentLoopRouter,
    RouterMetricsManager,
    SingletonThreadRouter,
)
from ray.serve._private.test_utils import FakeCounter, FakeGauge, MockTimer
from ray.serve._private.utils import (
    Semaphore,
    decompress_metric_report,
    get_random_string,
)
from ray.serve.config import AutoscalingConfig, RequestRouterConfig
from ray.serve.exceptions import (
    BackPressureError,
    DeploymentUnavailableError,
    ReplicaUnavailableError,
)


class FakeReplicaResult(ReplicaResult):
    def __init__(
        self,
        replica_id,
        is_generator_object: bool,
        queue_len_info: Optional[ReplicaQueueLengthInfo] = None,
    ):
        self._replica_id = replica_id
        self._is_generator_object = is_generator_object
        self._queue_len_info = queue_len_info
        self._done_callbacks: List[Callable] = []
        self.cancelled = False

    async def get_rejection_response(self):
        return self._queue_len_info

    def get(self, timeout_s: Optional[float]):
        raise NotImplementedError

    async def get_async(self):
        raise NotImplementedError

    def __next__(self):
        raise NotImplementedError

    async def __anext__(self):
        raise NotImplementedError

    def add_done_callback(self, callback: Callable):
        self._done_callbacks.append(callback)

    def fire_done_callbacks(self, result=None):
        """Simulate request completion by invoking all registered callbacks."""
        for cb in self._done_callbacks:
            cb(result)
        self._done_callbacks.clear()

    def cancel(self):
        self.cancelled = True

    def to_object_ref(self, timeout_s: Optional[float]) -> ray.ObjectRef:
        raise NotImplementedError

    async def to_object_ref_async(self) -> ray.ObjectRef:
        raise NotImplementedError

    def to_object_ref_gen(self) -> ray.ObjectRefGenerator:
        raise NotImplementedError


class FakeReplica(RunningReplica):
    def __init__(
        self,
        replica_id: ReplicaID,
        *,
        queue_len_info: Optional[ReplicaQueueLengthInfo] = None,
        is_cross_language: bool = False,
        error: Optional[Exception] = None,
        node_id: str = "fake-node-id",
        availability_zone: Optional[str] = None,
        node_ip: str = "127.0.0.1",
        port: Optional[int] = 8000,
        actor_id: Optional[ray.ActorID] = None,
    ):
        self._replica_id = replica_id
        self._is_cross_language = is_cross_language
        self._queue_len_info = queue_len_info
        self._error = error
        self._reserved_slots: Set[str] = set()
        self._slot_counter = 0
        self._requests_sent = []  # Track all requests sent to this replica
        self._reject_reservation = False
        self._reservation_queue_len = 1
        self._release_slot_error: Optional[Exception] = None

        # Create a minimal _replica_info object to satisfy router.py requirements
        self._replica_info = Mock()
        self._replica_info.node_id = node_id
        self._replica_info.availability_zone = availability_zone
        self._replica_info.node_ip = node_ip
        self._replica_info.port = port
        self._replica_info.replica_id = replica_id
        self._actor_id = actor_id

    @property
    def replica_id(self) -> ReplicaID:
        return self._replica_id

    @property
    def node_id(self) -> str:
        return self._replica_info.node_id

    @property
    def availability_zone(self) -> Optional[str]:
        return self._replica_info.availability_zone

    @property
    def actor_id(self) -> Optional[ray.ActorID]:
        return self._actor_id

    @property
    def is_cross_language(self) -> bool:
        return self._is_cross_language

    def get_queue_len(self, *, deadline_s: float) -> int:
        raise NotImplementedError

    async def reserve_slot(
        self, request_metadata: RequestMetadata
    ) -> Tuple[str, ReplicaQueueLengthInfo]:
        """Reserve a slot and return a token."""
        if self._reject_reservation:
            return "", ReplicaQueueLengthInfo(
                accepted=False,
                num_ongoing_requests=self._reservation_queue_len,
            )

        self._slot_counter += 1
        token = f"slot-token-{self._slot_counter}"
        self._reserved_slots.add(token)
        return token, ReplicaQueueLengthInfo(
            accepted=True,
            num_ongoing_requests=len(self._reserved_slots),
        )

    async def release_slot(self, slot_token: str) -> int:
        """Release a reserved slot."""
        if self._release_slot_error is not None:
            raise self._release_slot_error
        self._reserved_slots.discard(slot_token)
        return len(self._reserved_slots)

    def set_release_slot_error(self, error: Optional[Exception]) -> None:
        """Configure release_slot to raise the given error (or clear with None)."""
        self._release_slot_error = error

    def send_request_with_slot(
        self, pr: PendingRequest, slot_token: str
    ) -> FakeReplicaResult:
        """Send request using a reserved slot."""
        assert slot_token in self._reserved_slots, f"Invalid slot token: {slot_token}"

        # Create result same way as try_send_request
        if pr.metadata.is_streaming:
            return FakeReplicaResult(self._replica_id, is_generator_object=True)
        else:
            return FakeReplicaResult(self._replica_id, is_generator_object=False)

    def try_send_request(
        self, pr: PendingRequest, with_rejection: bool
    ) -> FakeReplicaResult:
        # Raise immediately if we're simulating send-time failure (e.g. ActorDiedError).
        # Real replicas can raise here when with_rejection=False (e.g. broadcast path).
        if self._error:
            raise self._error

        # Track the request
        if pr.metadata._reserved_slot_token:
            self._reserved_slots.discard(pr.metadata._reserved_slot_token)

        self._requests_sent.append(
            {
                "request_id": pr.metadata.request_id,
                "with_rejection": with_rejection,
            }
        )

        if with_rejection:
            assert (
                not self.is_cross_language
            ), "Rejection not supported for cross language."
            assert (
                self._queue_len_info is not None
            ), "Must set queue_len_info to use `send_request_with_rejection`."

            return FakeReplicaResult(
                self._replica_id,
                is_generator_object=True,
                queue_len_info=self._queue_len_info,
            )
        else:
            if pr.metadata.is_streaming:
                return FakeReplicaResult(self._replica_id, is_generator_object=True)
            else:
                return FakeReplicaResult(self._replica_id, is_generator_object=False)


class FakeRequestRouter(RequestRouter):
    def __init__(self, use_queue_len_cache: bool):
        self._block_requests = False
        self._blocked_requests: List[asyncio.Event] = []
        self._replica_to_return: Optional[FakeReplica] = None
        self._replica_to_return_on_retry: Optional[FakeReplica] = None
        self._replica_queue_len_cache = ReplicaQueueLengthCache()
        self._dropped_replicas: Set[ReplicaID] = set()
        self._use_queue_len_cache = use_queue_len_cache
        self._use_replica_queue_len_cache = use_queue_len_cache
        self.on_request_routed_called = False
        self.completed_requests: List[Tuple[ReplicaID, str]] = []

    def create_replica_wrapper(self, replica_info: RunningReplicaInfo):
        return FakeReplica(replica_info)

    @property
    def replica_queue_len_cache(self) -> ReplicaQueueLengthCache:
        return self._replica_queue_len_cache

    @property
    def dropped_replicas(self) -> Set[ReplicaID]:
        return self._dropped_replicas

    @property
    def curr_replicas(self) -> Dict[ReplicaID, RunningReplica]:
        replicas = {}
        if self._replica_to_return is not None:
            replicas[self._replica_to_return.replica_id] = self._replica_to_return
        if self._replica_to_return_on_retry is not None:
            replicas[
                self._replica_to_return_on_retry.replica_id
            ] = self._replica_to_return_on_retry

        return replicas

    def update_replicas(self, replicas: List[RunningReplica]):
        pass

    def on_replica_actor_died(self, replica_id: ReplicaID):
        self._dropped_replicas.add(replica_id)

    def on_new_queue_len_info(self, replica_id: ReplicaID, num_ongoing_requests: int):
        if self._use_queue_len_cache:
            self._replica_queue_len_cache.update(replica_id, num_ongoing_requests)

    def on_send_request(self, replica_id: ReplicaID):
        if self._use_queue_len_cache:
            num_ongoing_requests = self._replica_queue_len_cache.get(replica_id) or 0
            self._replica_queue_len_cache.update(replica_id, num_ongoing_requests + 1)

    def on_replica_actor_unavailable(self, replica_id: ReplicaID):
        self._replica_queue_len_cache.invalidate_key(replica_id)

    def set_should_block_requests(self, block_requests: bool):
        self._block_requests = block_requests

    def set_replica_to_return(self, replica: FakeReplica):
        self._replica_to_return = replica

    def set_replica_to_return_on_retry(self, replica: FakeReplica):
        self._replica_to_return_on_retry = replica

    def unblock_requests(self, num: int):
        assert self._block_requests and len(self._blocked_requests) >= num
        for _ in range(num):
            self._blocked_requests.pop(0).set()

    async def _choose_replica_for_request(
        self, pr: PendingRequest, *, is_retry: bool = False
    ) -> FakeReplica:
        if self._block_requests:
            event = asyncio.Event()
            self._blocked_requests.append(event)
            await event.wait()

        if is_retry:
            assert (
                self._replica_to_return_on_retry is not None
            ), "Set a replica to return on retry."
            replica = self._replica_to_return_on_retry
        else:
            assert self._replica_to_return is not None, "Set a replica to return."
            replica = self._replica_to_return

        return replica

    async def choose_replicas(
        self,
        candidate_replicas: List[RunningReplica],
        pending_request: Optional[PendingRequest] = None,
    ) -> List[List[RunningReplica]]:
        pass

    def on_request_routed(
        self,
        pending_request: PendingRequest,
        replica_id: ReplicaID,
        result: ReplicaResult,
    ) -> None:
        self.on_request_routed_called = True

    def on_request_completed(
        self,
        replica_id: ReplicaID,
        internal_request_id: str,
    ) -> None:
        self.completed_requests.append((replica_id, internal_request_id))


@pytest.fixture
async def setup_router(request) -> Tuple[AsyncioRouter, FakeRequestRouter]:
    if not hasattr(request, "param"):
        request.param = {}

    fake_request_router = FakeRequestRouter(
        request.param.get("enable_queue_len_cache", False)
    )
    router = AsyncioRouter(
        # TODO(edoakes): refactor to make a better fake controller or not depend on it.
        controller_handle=Mock(),
        deployment_id=DeploymentID(name="test-deployment"),
        handle_id="test-handle-id",
        self_actor_id="test-node-id",
        handle_source=DeploymentHandleSource.UNKNOWN,
        event_loop=get_or_create_event_loop(),
        enable_strict_max_ongoing_requests=request.param.get(
            "enable_strict_max_ongoing_requests", False
        ),
        request_router=fake_request_router,
        node_id="test-node-id",
        availability_zone="test-az",
        prefer_local_node_routing=False,
        _request_router_initialized_event=asyncio.Event(),
    )
    yield router, fake_request_router
    await router.shutdown()


def dummy_request_metadata(is_streaming: bool = False) -> RequestMetadata:
    return RequestMetadata(
        request_id="test-request-1",
        internal_request_id="test-internal-request-1",
        is_streaming=is_streaming,
    )


class FakeReplicaMetricsManager:
    def __init__(self):
        self.num_ongoing_requests = 0

    def get_num_ongoing_requests(self) -> int:
        return self.num_ongoing_requests

    def inc_num_ongoing_requests(self, request_metadata: RequestMetadata):
        self.num_ongoing_requests += 1

    def dec_num_ongoing_requests(self, request_metadata: RequestMetadata):
        self.num_ongoing_requests -= 1


class FakeServeReplicaForSlotReservation(ServeReplica):
    def __init__(
        self,
        *,
        max_ongoing_requests: int = 1,
        graceful_shutdown_wait_loop_s: float = 0.01,
    ):
        self._deployment_config = Mock(
            max_ongoing_requests=max_ongoing_requests,
            graceful_shutdown_wait_loop_s=graceful_shutdown_wait_loop_s,
        )
        self._metrics_manager = FakeReplicaMetricsManager()
        self._reserved_slots = set()
        self._semaphore = Semaphore(lambda: self.max_ongoing_requests)
        # The real __init__ is bypassed; set the quiesce flag read by
        # _can_accept_request (reservations are rejected once quiescing).
        self._quiescing = False


@pytest.mark.asyncio
class TestReplicaSlotReservation:
    async def test_reserved_slot_counts_against_replica_capacity(self):
        # A reserved slot consumes capacity, so another reservation should be
        # rejected once the replica reaches max_ongoing_requests.
        replica = FakeServeReplicaForSlotReservation()

        request_metadata = dummy_request_metadata()

        slot_token = "slot-token-1"
        accepted, queue_len = await replica.reserve_slot(request_metadata, slot_token)
        assert accepted
        assert slot_token in replica._reserved_slots
        assert queue_len == 1
        assert replica.get_num_ongoing_requests() == 1

        accepted, queue_len = await replica.reserve_slot(
            request_metadata, "slot-token-2"
        )
        assert not accepted
        assert queue_len == 1

        dispatch_metadata = replace(request_metadata, _reserved_slot_token=slot_token)
        async with replica._start_request(dispatch_metadata):
            assert slot_token not in replica._reserved_slots
            assert replica.get_num_ongoing_requests() == 1

        assert replica.get_num_ongoing_requests() == 0

    async def test_release_slot_returns_status_and_recovers_capacity(self):
        # Releasing a known token reports True and frees a slot; releasing an
        # unknown or already-released token reports False and is a no-op.
        replica = FakeServeReplicaForSlotReservation()

        request_metadata = dummy_request_metadata()
        accepted, _ = await replica.reserve_slot(request_metadata, "tok-a")
        assert accepted
        assert replica.get_num_ongoing_requests() == 1

        accepted, count = replica.release_slot("tok-a")
        assert accepted
        assert count == 0
        assert replica.get_num_ongoing_requests() == 0

        # Releasing the same token a second time is a no-op.
        accepted, count = replica.release_slot("tok-a")
        assert not accepted
        assert count == 0

        # Releasing a token that was never reserved is also a no-op.
        accepted, count = replica.release_slot("never-reserved")
        assert not accepted
        assert count == 0

        # Capacity is fully recovered: a fresh reservation succeeds.
        accepted, _ = await replica.reserve_slot(request_metadata, "tok-b")
        assert accepted

    async def test_dispatch_with_unknown_slot_token_is_rejected(self):
        # A request that arrives with a slot token nobody reserved (or that
        # was already released) must be rejected loudly and must not
        # consume capacity.
        replica = FakeServeReplicaForSlotReservation()

        bogus = replace(dummy_request_metadata(), _reserved_slot_token="never-reserved")
        with pytest.raises(RuntimeError, match="unknown reserved slot"):
            async with replica._start_request(bogus):
                pass

        assert replica.get_num_ongoing_requests() == 0

        # The replica is still healthy: a normal request can still run.
        async with replica._start_request(dummy_request_metadata()):
            assert replica.get_num_ongoing_requests() == 1
        assert replica.get_num_ongoing_requests() == 0

    async def test_drain_waits_for_reserved_slots(self):
        # A reserved-but-not-yet-dispatched slot is holding capacity, so a
        # graceful shutdown must wait for it just like an in-flight request.
        # Otherwise the replica can finish draining between reserve_slot and
        # the dispatch, leaving the router to send the request to a dead replica.
        replica = FakeServeReplicaForSlotReservation(graceful_shutdown_wait_loop_s=0.01)

        request_metadata = dummy_request_metadata()
        accepted, _ = await replica.reserve_slot(request_metadata, "slot-token-1")
        assert accepted

        drain_task = asyncio.create_task(replica._drain_ongoing_requests())
        try:
            await asyncio.sleep(0.1)
            assert (
                not drain_task.done()
            ), "drain should still be running while a slot is reserved"

            replica.release_slot("slot-token-1")
            await asyncio.wait_for(drain_task, timeout=1.0)
        finally:
            if not drain_task.done():
                drain_task.cancel()

    async def test_drain_honors_min_period_then_waits_for_ongoing(self):
        # The minimum draining period keeps the replica alive (still serving) so
        # external load balancers deregister it. The drain must not exit until the
        # period has elapsed AND there are no ongoing requests -- so a request
        # admitted during the period and still running keeps it alive, instead of
        # being reset when the replica tears down.
        replica = FakeServeReplicaForSlotReservation(graceful_shutdown_wait_loop_s=0.02)
        min_period_s = 0.2

        drain_task = asyncio.create_task(
            replica._drain_ongoing_requests(min_draining_period_s=min_period_s)
        )
        try:
            # Nothing in flight, but the minimum period has not elapsed yet.
            await asyncio.sleep(0.1)
            assert not drain_task.done(), "drain must honor the minimum period"

            # A request becomes ongoing partway through the period.
            replica._metrics_manager.num_ongoing_requests = 1

            # The period elapses, but the request is still ongoing -> keep waiting.
            await asyncio.sleep(0.2)
            assert (
                not drain_task.done()
            ), "drain must not exit while a request is still ongoing"

            # Request finishes; only now may the drain complete.
            replica._metrics_manager.num_ongoing_requests = 0
            await asyncio.wait_for(drain_task, timeout=1.0)
        finally:
            if not drain_task.done():
                drain_task.cancel()

    async def test_direct_ingress_request_is_rejected(self):
        replica = FakeServeReplicaForSlotReservation()

        request_metadata = replace(dummy_request_metadata(), is_direct_ingress=True)

        with pytest.raises(RuntimeError, match="direct-ingress"):
            await replica.reserve_slot(request_metadata, "slot-token-1")

        assert replica.get_num_ongoing_requests() == 0


class FakeRunningReplicaForSlotReservation(RunningReplica):
    def __init__(self, actor_handle: Mock):
        self._replica_info = Mock(is_cross_language=False)
        self._actor_handle = actor_handle


@pytest.mark.asyncio
class TestRunningReplicaSlotReservation:
    async def test_cancellation_releases_reserved_slot(self):
        # Cancellation can arrive after the actor has already reserved the slot,
        # so ray.cancel(obj_ref) is a no-op and release_slot.remote() is the
        # only one that frees the leaked slot.
        pending = asyncio.get_running_loop().create_future()
        actor_handle = Mock()
        actor_handle.reserve_slot.remote = Mock(return_value=pending)
        actor_handle.release_slot.remote = Mock()
        replica = FakeRunningReplicaForSlotReservation(actor_handle)

        with patch("ray.cancel") as mock_ray_cancel:
            task = asyncio.create_task(replica.reserve_slot(dummy_request_metadata()))
            # Yield so the task reaches `await obj_ref` before we cancel.
            await asyncio.sleep(0)
            task.cancel()

            with pytest.raises(asyncio.CancelledError):
                await task

        mock_ray_cancel.assert_called_once_with(pending)

        actor_handle.release_slot.remote.assert_called_once()
        released_token = actor_handle.release_slot.remote.call_args[0][0]
        reserved_token = actor_handle.reserve_slot.remote.call_args[0][1]
        assert released_token == reserved_token

    @pytest.mark.parametrize(
        "exc",
        [
            ActorUnavailableError(error_message="unavailable", actor_id=None),
            ActorDiedError(),
            RuntimeError("boom"),
        ],
    )
    async def test_exception_releases_reserved_slot(self, exc):
        # If reserve_slot's reply fails (e.g. ActorUnavailableError),
        # the actor may have already reserved the slot. Best-effort
        # release_slot.remote() prevents the leak.
        pending = asyncio.get_running_loop().create_future()
        pending.set_exception(exc)
        actor_handle = Mock()
        actor_handle.reserve_slot.remote = Mock(return_value=pending)
        actor_handle.release_slot.remote = Mock()
        replica = FakeRunningReplicaForSlotReservation(actor_handle)

        with pytest.raises(type(exc)):
            await replica.reserve_slot(dummy_request_metadata())

        actor_handle.release_slot.remote.assert_called_once()
        released_token = actor_handle.release_slot.remote.call_args[0][0]
        reserved_token = actor_handle.reserve_slot.remote.call_args[0][1]
        assert released_token == reserved_token


@pytest.mark.asyncio
class TestBroadcast:
    async def test_unavailable_fails_without_waiting_for_router_init(self):
        fake_request_router = FakeRequestRouter(use_queue_len_cache=False)
        router = AsyncioRouter(
            controller_handle=Mock(),
            deployment_id=DeploymentID(name="test-deployment"),
            handle_id="test-handle-id",
            self_actor_id="test-node-id",
            handle_source=DeploymentHandleSource.UNKNOWN,
            event_loop=get_or_create_event_loop(),
            enable_strict_max_ongoing_requests=False,
            request_router=fake_request_router,
            node_id="test-node-id",
            availability_zone="test-az",
            prefer_local_node_routing=False,
            _request_router_initialized_event=asyncio.Event(),
        )
        # Simulate a router that has not finished initialization yet.
        router._request_router_initialized.clear()
        router._deployment_available = False

        with pytest.raises(DeploymentUnavailableError):
            await asyncio.wait_for(
                router.broadcast(dummy_request_metadata()),
                timeout=0.1,
            )

    async def test_skips_dead_replica_and_continues(self, setup_router):
        router, fake_request_router = setup_router
        d_id = DeploymentID(name="test")
        r1_id = ReplicaID(unique_id="r1", deployment_id=d_id)
        r2_id = ReplicaID(unique_id="r2", deployment_id=d_id)

        fake_request_router.set_replica_to_return(
            FakeReplica(r1_id, error=ActorDiedError())
        )
        fake_request_router.set_replica_to_return_on_retry(FakeReplica(r2_id))

        results = await router.broadcast(dummy_request_metadata())
        assert len(results) == 1
        assert results[0]._replica_id == r2_id
        assert r1_id in fake_request_router.dropped_replicas

    async def test_all_replicas_dead_raises_unavailable(self, setup_router):
        router, fake_request_router = setup_router
        d_id = DeploymentID(name="test")
        r1_id = ReplicaID(unique_id="r1", deployment_id=d_id)

        fake_request_router.set_replica_to_return(
            FakeReplica(r1_id, error=ActorDiedError())
        )

        with pytest.raises(DeploymentUnavailableError):
            await router.broadcast(dummy_request_metadata())


class TestDeploymentBroadcastResponseLocalTesting:
    """Tests for DeploymentBroadcastResponse with loop=None (local testing mode)."""

    def test_results_twice_after_error_does_not_raise_coroutine_reuse(self):
        """Calling results() twice when the coroutine raises should replay the
        cached exception, not crash with 'cannot reuse already awaited coroutine'.
        """
        from ray.serve.handle import DeploymentBroadcastResponse

        async def failing_coro():
            raise DeploymentUnavailableError("no replicas")

        response = DeploymentBroadcastResponse(failing_coro(), loop=None)

        with pytest.raises(DeploymentUnavailableError):
            response.results()

        # Second call must raise the same error, not RuntimeError about reuse.
        with pytest.raises(DeploymentUnavailableError):
            response.results()

    def test_results_twice_after_success_returns_same_results(self):
        """Calling results() twice on a successful broadcast should return the
        same cached results both times.
        """
        from ray.serve.handle import DeploymentBroadcastResponse

        fake_result = FakeReplicaResult(
            replica_id=ReplicaID(
                unique_id="r1", deployment_id=DeploymentID(name="test")
            ),
            is_generator_object=False,
        )

        async def success_coro():
            return [fake_result]

        response = DeploymentBroadcastResponse(success_coro(), loop=None)

        # First call succeeds and caches.
        replica_results = response._fetch_replica_results_sync()
        assert len(replica_results) == 1

        # Second call returns the same cached list.
        replica_results_2 = response._fetch_replica_results_sync()
        assert replica_results_2 is replica_results


class FakeGrpcCallForRejection:
    """Minimal stand-in for a grpc.aio.Call whose connection attempt fails."""

    def __init__(self, error: Exception):
        self._error = error

    async def wait_for_connection(self):
        raise self._error


class FakeGrpcReplicaResultForRejection(gRPCReplicaResult):
    """Bypasses the real __init__ (context bookkeeping, background tasks);
    sets only the attributes get_rejection_response() reads."""

    def __init__(self, call, *, is_streaming: bool = False):
        self._call = call
        self._with_rejection = True
        self._rejection_response = None
        self._is_streaming = is_streaming
        # ActorUnavailableError -> ActorID(actor_id) requires a valid 16-byte
        # binary, so use a real random ActorID rather than an arbitrary value.
        self._actor_id = ray.ActorID.from_random()


def _aio_rpc_error(
    code: grpc.StatusCode, initial_metadata: Optional[grpc.aio.Metadata] = None
) -> grpc.aio.AioRpcError:
    return grpc.aio.AioRpcError(
        code=code,
        initial_metadata=initial_metadata or grpc.aio.Metadata(),
        trailing_metadata=grpc.aio.Metadata(),
        details=code.name,
        debug_error_string=None,
    )


@pytest.mark.asyncio
class TestGrpcRejectionResponseErrorMapping:
    """get_rejection_response() error mapping for dispatches that provably
    never started executing (no `accepted` initial metadata received).

    CANCELLED is what a replica's gRPC server returns for streams landing in
    its `server.stop()` window during graceful shutdown (quiesce). Like
    UNAVAILABLE, it must surface as ActorUnavailableError so the router
    retries the request on another replica instead of failing it.
    """

    @pytest.mark.parametrize(
        "code", [grpc.StatusCode.CANCELLED, grpc.StatusCode.UNAVAILABLE]
    )
    async def test_pre_accept_errors_map_to_actor_unavailable(self, code):
        result = FakeGrpcReplicaResultForRejection(
            FakeGrpcCallForRejection(_aio_rpc_error(code))
        )

        with pytest.raises(ActorUnavailableError):
            await result.get_rejection_response()

    async def test_other_codes_propagate_raw(self):
        result = FakeGrpcReplicaResultForRejection(
            FakeGrpcCallForRejection(_aio_rpc_error(grpc.StatusCode.INTERNAL))
        )

        with pytest.raises(grpc.aio.AioRpcError):
            await result.get_rejection_response()

    async def test_accepted_metadata_on_error_is_not_retried(self):
        # A unary call that fails mid-execution carries `accepted=1` initial
        # metadata in the error: the request DID start executing, so it must
        # NOT be converted into a retryable ActorUnavailableError.
        error = _aio_rpc_error(
            grpc.StatusCode.CANCELLED,
            initial_metadata=grpc.aio.Metadata(
                ("accepted", "1"), ("num_ongoing_requests", "3")
            ),
        )
        result = FakeGrpcReplicaResultForRejection(FakeGrpcCallForRejection(error))

        info = await result.get_rejection_response()
        assert info.accepted is True
        assert info.num_ongoing_requests == 3

    async def test_streaming_pre_accept_cancelled_maps_to_actor_unavailable(self):
        # For streaming calls initial metadata is sent before execution, so
        # an error during connection establishment always means the request
        # never started; the unary metadata-recovery branch is skipped.
        error = _aio_rpc_error(
            grpc.StatusCode.CANCELLED,
            initial_metadata=grpc.aio.Metadata(
                ("accepted", "1"), ("num_ongoing_requests", "3")
            ),
        )
        result = FakeGrpcReplicaResultForRejection(
            FakeGrpcCallForRejection(error), is_streaming=True
        )

        with pytest.raises(ActorUnavailableError):
            await result.get_rejection_response()


@pytest.mark.asyncio
class TestAssignRequest:
    @pytest.mark.parametrize("is_streaming", [False, True])
    async def test_basic(
        self,
        setup_router: Tuple[AsyncioRouter, FakeRequestRouter],
        is_streaming: bool,
    ):
        router, fake_request_router = setup_router

        r1_id = ReplicaID(
            unique_id="test-replica-1", deployment_id=DeploymentID(name="test")
        )
        replica = FakeReplica(r1_id)
        fake_request_router.set_replica_to_return(replica)

        request_metadata = RequestMetadata(
            request_id="test-request-1",
            internal_request_id="test-internal-request-1",
            is_streaming=is_streaming,
        )
        replica_result = await router.assign_request(request_metadata)
        if is_streaming:
            assert replica_result._is_generator_object
            assert replica_result._replica_id == r1_id
        else:
            assert not replica_result._is_generator_object
            assert replica_result._replica_id == r1_id

    @pytest.mark.parametrize(
        "setup_router",
        [
            {
                "enable_strict_max_ongoing_requests": True,
                "enable_queue_len_cache": False,
            },
            {
                "enable_strict_max_ongoing_requests": True,
                "enable_queue_len_cache": True,
            },
        ],
        indirect=True,
    )
    @pytest.mark.parametrize("is_streaming", [False, True])
    async def test_basic_with_rejection(
        self,
        setup_router: Tuple[AsyncioRouter, FakeRequestRouter],
        is_streaming: bool,
    ):
        router, fake_request_router = setup_router

        r1_id = ReplicaID(
            unique_id="test-replica-1", deployment_id=DeploymentID(name="test")
        )
        replica = FakeReplica(
            r1_id,
            queue_len_info=ReplicaQueueLengthInfo(
                accepted=True, num_ongoing_requests=10
            ),
        )
        fake_request_router.set_replica_to_return(replica)

        request_metadata = RequestMetadata(
            request_id="test-request-1",
            internal_request_id="test-internal-request-1",
            is_streaming=is_streaming,
        )
        replica_result = await router.assign_request(request_metadata)
        assert replica_result._is_generator_object
        assert replica_result._replica_id == r1_id

        if router._request_router._use_queue_len_cache:
            assert fake_request_router.replica_queue_len_cache.get(r1_id) == 10

    @pytest.mark.parametrize(
        "setup_router",
        [
            {
                "enable_strict_max_ongoing_requests": True,
                "enable_queue_len_cache": False,
            },
            {
                "enable_strict_max_ongoing_requests": True,
                "enable_queue_len_cache": True,
            },
        ],
        indirect=True,
    )
    @pytest.mark.parametrize("is_streaming", [False, True])
    async def test_retry_with_rejection(
        self,
        setup_router: Tuple[AsyncioRouter, FakeRequestRouter],
        is_streaming: bool,
    ):
        router, fake_request_router = setup_router

        r1_id = ReplicaID(
            unique_id="test-replica-1", deployment_id=DeploymentID(name="test")
        )
        replica1 = FakeReplica(
            r1_id,
            queue_len_info=ReplicaQueueLengthInfo(
                accepted=False, num_ongoing_requests=10
            ),
        )
        fake_request_router.set_replica_to_return(replica1)

        r2_id = ReplicaID(
            unique_id="test-replica-2", deployment_id=DeploymentID(name="test")
        )
        replica2 = FakeReplica(
            r2_id,
            queue_len_info=ReplicaQueueLengthInfo(
                accepted=True, num_ongoing_requests=20
            ),
        )
        fake_request_router.set_replica_to_return_on_retry(replica2)

        request_metadata = RequestMetadata(
            request_id="test-request-1",
            internal_request_id="test-internal-request-1",
            is_streaming=is_streaming,
        )
        replica_result = await router.assign_request(request_metadata)
        assert replica_result._is_generator_object
        assert replica_result._replica_id == r2_id

        if router._request_router._use_queue_len_cache:
            assert fake_request_router.replica_queue_len_cache.get(r1_id) == 10
            assert fake_request_router.replica_queue_len_cache.get(r2_id) == 20

    @pytest.mark.parametrize(
        "setup_router",
        [{"enable_strict_max_ongoing_requests": True}],
        indirect=True,
    )
    async def test_cross_lang_no_rejection(
        self, setup_router: Tuple[AsyncioRouter, FakeRequestRouter]
    ):
        router, fake_request_router = setup_router

        r1_id = ReplicaID(
            unique_id="test-replica-1", deployment_id=DeploymentID(name="test")
        )
        replica = FakeReplica(r1_id, is_cross_language=True)
        fake_request_router.set_replica_to_return(replica)

        request_metadata = RequestMetadata(
            request_id="test-request-1",
            internal_request_id="test-internal-request-1",
        )
        replica_result = await router.assign_request(request_metadata)
        assert not replica_result._is_generator_object
        assert replica_result._replica_id == r1_id

    async def test_max_queued_requests_no_limit(
        self, setup_router: Tuple[AsyncioRouter, FakeRequestRouter]
    ):
        router, fake_request_router = setup_router
        fake_request_router.set_should_block_requests(True)
        router.update_deployment_config(DeploymentConfig(max_queued_requests=-1))

        r1_id = ReplicaID(
            unique_id="test-replica-1", deployment_id=DeploymentID(name="test")
        )
        replica = FakeReplica(r1_id)
        fake_request_router.set_replica_to_return(replica)

        request_metadata = RequestMetadata(
            request_id="test-request-1",
            internal_request_id="test-internal-request-1",
        )

        # Queued a bunch of tasks. None should error because there's no limit.
        assign_request_tasks = [
            asyncio.ensure_future(router.assign_request(request_metadata))
            for _ in range(100)
        ]

        _, pending = await asyncio.wait(assign_request_tasks, timeout=0.01)
        assert len(pending) == len(assign_request_tasks)

        # Unblock the requests, now they should all get routed.
        fake_request_router.unblock_requests(100)
        assert all(
            [
                not replica_result._is_generator_object
                and replica_result._replica_id == r1_id
                for replica_result in await asyncio.gather(*assign_request_tasks)
            ]
        )

    async def test_max_queued_requests_limited(
        self, setup_router: Tuple[AsyncioRouter, FakeRequestRouter]
    ):
        router, fake_request_router = setup_router
        fake_request_router.set_should_block_requests(True)
        router.update_deployment_config(DeploymentConfig(max_queued_requests=5))

        r1_id = ReplicaID(
            unique_id="test-replica-1", deployment_id=DeploymentID(name="test")
        )
        replica = FakeReplica(r1_id)
        fake_request_router.set_replica_to_return(replica)

        request_metadata = RequestMetadata(
            request_id="test-request-1",
            internal_request_id="test-internal-request-1",
        )

        # Queued `max_queued_requests` tasks. None should fail.
        assign_request_tasks = [
            asyncio.ensure_future(router.assign_request(request_metadata))
            for _ in range(5)
        ]

        _, pending = await asyncio.wait(assign_request_tasks, timeout=0.01)
        assert len(pending) == len(assign_request_tasks)

        # Try to queue more tasks, they should fail immediately.
        for _ in range(10):
            with pytest.raises(BackPressureError):
                await router.assign_request(request_metadata)

        # Unblock a request.
        fake_request_router.unblock_requests(1)
        done, pending = await asyncio.wait(assign_request_tasks, timeout=0.01)
        assert len(done) == 1
        replica_result = await done.pop()
        assert not replica_result._is_generator_object
        assert replica_result._replica_id == r1_id

        # One more task should be allowed to be queued.
        assign_request_tasks = list(pending) + [
            asyncio.ensure_future(router.assign_request(request_metadata))
        ]

        _, pending = await asyncio.wait(assign_request_tasks, timeout=0.01)
        assert len(pending) == len(assign_request_tasks)

        # Unblock the requests, now they should all get routed.
        fake_request_router.unblock_requests(5)
        assert all(
            [
                not replica_result._is_generator_object
                and replica_result._replica_id == r1_id
                for replica_result in await asyncio.gather(*assign_request_tasks)
            ]
        )

    async def test_max_queued_requests_updated(
        self, setup_router: Tuple[AsyncioRouter, FakeRequestRouter]
    ):
        router, fake_request_router = setup_router
        fake_request_router.set_should_block_requests(True)
        router.update_deployment_config(DeploymentConfig(max_queued_requests=5))

        r1_id = ReplicaID(
            unique_id="test-replica-1", deployment_id=DeploymentID(name="test")
        )
        replica = FakeReplica(r1_id)
        fake_request_router.set_replica_to_return(replica)

        request_metadata = RequestMetadata(
            request_id="test-request-1",
            internal_request_id="test-internal-request-1",
        )

        # Queued `max_queued_requests` tasks. None should fail.
        assign_request_tasks = [
            asyncio.ensure_future(router.assign_request(request_metadata))
            for _ in range(5)
        ]

        _, pending = await asyncio.wait(assign_request_tasks, timeout=0.01)
        assert len(pending) == len(assign_request_tasks)

        # Try to queue more tasks, they should fail immediately.
        for _ in range(10):
            with pytest.raises(BackPressureError):
                await router.assign_request(request_metadata)

        # Dynamically increase `max_queued_requests`, more tasks should be allowed to
        # be queued.
        router.update_deployment_config(DeploymentConfig(max_queued_requests=10))
        assign_request_tasks.extend(
            [
                asyncio.ensure_future(router.assign_request(request_metadata))
                for _ in range(5)
            ]
        )

        _, pending = await asyncio.wait(assign_request_tasks, timeout=0.01)
        assert len(pending) == len(assign_request_tasks)

        # Try to queue more tasks, they should fail immediately.
        for _ in range(10):
            with pytest.raises(BackPressureError):
                await router.assign_request(request_metadata)

        # Dynamically decrease `max_queued_requests`, the existing tasks should remain
        # queued but the limit should apply to new tasks.
        router.update_deployment_config(DeploymentConfig(max_queued_requests=5))
        _, pending = await asyncio.wait(assign_request_tasks, timeout=0.01)
        assert len(pending) == len(assign_request_tasks)

        for _ in range(10):
            with pytest.raises(BackPressureError):
                await router.assign_request(request_metadata)

        fake_request_router.unblock_requests(5)
        done, pending = await asyncio.wait(assign_request_tasks, timeout=0.01)
        assert len(pending) == 5
        assert all(
            [
                not replica_result._is_generator_object
                and replica_result._replica_id == r1_id
                for replica_result in await asyncio.gather(*done)
            ]
        )
        assign_request_tasks = list(pending)

        # Try to queue more tasks, they should fail immediately if the new limit is
        # respected.
        for _ in range(10):
            with pytest.raises(BackPressureError):
                await router.assign_request(request_metadata)

        # Unblock the requests, now they should all get routed.
        fake_request_router.unblock_requests(5)
        assert all(
            [
                not replica_result._is_generator_object
                and replica_result._replica_id == r1_id
                for replica_result in await asyncio.gather(*assign_request_tasks)
            ]
        )

    @pytest.mark.parametrize(
        "setup_router",
        [
            {
                "enable_strict_max_ongoing_requests": True,
                "enable_queue_len_cache": True,
            },
        ],
        indirect=True,
    )
    async def test_replica_actor_died(
        self, setup_router: Tuple[AsyncioRouter, FakeRequestRouter]
    ):
        """ActorDiedError whose actor_id matches the replica → replica is dropped."""
        router, fake_request_router = setup_router
        d_id = DeploymentID(name="test")
        r1_id = ReplicaID(unique_id="r1", deployment_id=d_id)
        r2_id = ReplicaID(unique_id="r2", deployment_id=d_id)

        r1_actor_id = ray.ActorID.from_random()
        error = ActorDiedError()
        error.actor_id = r1_actor_id.hex()

        fake_request_router.set_replica_to_return(
            FakeReplica(r1_id, error=error, actor_id=r1_actor_id)
        )
        fake_request_router.set_replica_to_return_on_retry(
            FakeReplica(
                r2_id,
                queue_len_info=ReplicaQueueLengthInfo(
                    accepted=True, num_ongoing_requests=5
                ),
            )
        )
        await router.assign_request(dummy_request_metadata())
        assert r1_id in fake_request_router.dropped_replicas

    @pytest.mark.parametrize(
        "setup_router",
        [
            {
                "enable_strict_max_ongoing_requests": True,
                "enable_queue_len_cache": True,
            },
        ],
        indirect=True,
    )
    async def test_replica_actor_died_mismatched_actor_id(
        self, setup_router: Tuple[AsyncioRouter, FakeRequestRouter]
    ):
        """ActorDiedError whose actor_id does NOT match the replica → replica is NOT
        dropped (the error came from an upstream dependency, not this replica)."""
        router, fake_request_router = setup_router
        d_id = DeploymentID(name="test")
        r1_id = ReplicaID(unique_id="r1", deployment_id=d_id)
        r2_id = ReplicaID(unique_id="r2", deployment_id=d_id)

        # r1 has one actor_id but the error reports a different actor_id (upstream).
        r1_actor_id = ray.ActorID.from_random()
        upstream_actor_id = ray.ActorID.from_random()
        error = ActorDiedError()
        error.actor_id = upstream_actor_id.hex()

        fake_request_router.set_replica_to_return(
            FakeReplica(r1_id, error=error, actor_id=r1_actor_id)
        )
        # r2 handles the retry after r1 fails.
        fake_request_router.set_replica_to_return_on_retry(
            FakeReplica(
                r2_id,
                queue_len_info=ReplicaQueueLengthInfo(
                    accepted=True, num_ongoing_requests=5
                ),
            )
        )
        await router.assign_request(dummy_request_metadata())
        # r1 should NOT be dropped because the error came from a different actor.
        assert r1_id not in fake_request_router.dropped_replicas

    async def test_actor_died_matching_id_via_callback(
        self, setup_router: Tuple[AsyncioRouter, FakeRequestRouter]
    ):
        """Done-callback with ActorDiedError whose actor_id matches the replica →
        replica is dropped via _process_finished_request."""
        router, fake_request_router = setup_router
        d_id = DeploymentID(name="test")
        r1_id = ReplicaID(unique_id="r1", deployment_id=d_id)

        r1_actor_id = ray.ActorID.from_random()
        fake_request_router.set_replica_to_return(
            FakeReplica(r1_id, actor_id=r1_actor_id)
        )

        result = await router.assign_request(dummy_request_metadata())

        # Simulate the response finishing with an ActorDiedError from this replica.
        error = ActorDiedError()
        error.actor_id = r1_actor_id.hex()
        result.fire_done_callbacks(result=error)
        await asyncio.sleep(0)

        assert r1_id in fake_request_router.dropped_replicas

    async def test_actor_died_mismatched_id_via_callback(
        self, setup_router: Tuple[AsyncioRouter, FakeRequestRouter]
    ):
        """Done-callback with ActorDiedError whose actor_id does NOT match the
        replica → replica is NOT dropped (upstream dependency failure)."""
        router, fake_request_router = setup_router
        d_id = DeploymentID(name="test")
        r1_id = ReplicaID(unique_id="r1", deployment_id=d_id)

        r1_actor_id = ray.ActorID.from_random()
        upstream_actor_id = ray.ActorID.from_random()
        fake_request_router.set_replica_to_return(
            FakeReplica(r1_id, actor_id=r1_actor_id)
        )

        result = await router.assign_request(dummy_request_metadata())

        # Simulate the response finishing with an ActorDiedError from a different actor.
        error = ActorDiedError()
        error.actor_id = upstream_actor_id.hex()
        result.fire_done_callbacks(result=error)
        await asyncio.sleep(0)

        assert r1_id not in fake_request_router.dropped_replicas

    async def test_actor_died_none_actor_id_via_callback(
        self, setup_router: Tuple[AsyncioRouter, FakeRequestRouter]
    ):
        """Done-callback with ActorDiedError(actor_id=None) and replica has valid
        actor_id → replica IS dropped (conservative fallback when ID unknown)."""
        router, fake_request_router = setup_router
        d_id = DeploymentID(name="test")
        r1_id = ReplicaID(unique_id="r1", deployment_id=d_id)

        r1_actor_id = ray.ActorID.from_random()
        fake_request_router.set_replica_to_return(
            FakeReplica(r1_id, actor_id=r1_actor_id)
        )

        result = await router.assign_request(dummy_request_metadata())

        # ActorDiedError without actor_id (e.g., constructed without
        # ActorDiedErrorContext). Conservative fallback: mark replica dead.
        error = ActorDiedError()
        error.actor_id = None  # Explicitly None
        result.fire_done_callbacks(result=error)
        await asyncio.sleep(0)

        assert r1_id in fake_request_router.dropped_replicas

    async def test_ray_task_error_wrapping_actor_died_matching_id_via_callback(
        self, setup_router: Tuple[AsyncioRouter, FakeRequestRouter]
    ):
        """Done-callback with RayTaskError(cause=ActorDiedError) where actor_id
        matches the replica → replica is dropped."""
        router, fake_request_router = setup_router
        d_id = DeploymentID(name="test")
        r1_id = ReplicaID(unique_id="r1", deployment_id=d_id)

        r1_actor_id = ray.ActorID.from_random()
        fake_request_router.set_replica_to_return(
            FakeReplica(r1_id, actor_id=r1_actor_id)
        )

        result = await router.assign_request(dummy_request_metadata())

        actor_died = ActorDiedError()
        actor_died.actor_id = r1_actor_id.hex()
        wrapped = RayTaskError(
            function_name="test_func",
            traceback_str="",
            cause=actor_died,
            proctitle="test",
            pid=12345,
            ip="127.0.0.1",
        )
        result.fire_done_callbacks(result=wrapped)
        await asyncio.sleep(0)

        assert r1_id in fake_request_router.dropped_replicas

    async def test_ray_task_error_wrapping_actor_died_mismatched_id_via_callback(
        self, setup_router: Tuple[AsyncioRouter, FakeRequestRouter]
    ):
        """Done-callback with RayTaskError(cause=ActorDiedError) where actor_id does
        NOT match the replica → replica is NOT dropped (upstream failure)."""
        router, fake_request_router = setup_router
        d_id = DeploymentID(name="test")
        r1_id = ReplicaID(unique_id="r1", deployment_id=d_id)

        r1_actor_id = ray.ActorID.from_random()
        upstream_actor_id = ray.ActorID.from_random()
        fake_request_router.set_replica_to_return(
            FakeReplica(r1_id, actor_id=r1_actor_id)
        )

        result = await router.assign_request(dummy_request_metadata())

        actor_died = ActorDiedError()
        actor_died.actor_id = upstream_actor_id.hex()
        wrapped = RayTaskError(
            function_name="test_func",
            traceback_str="",
            cause=actor_died,
            proctitle="test",
            pid=12345,
            ip="127.0.0.1",
        )
        result.fire_done_callbacks(result=wrapped)
        await asyncio.sleep(0)

        assert r1_id not in fake_request_router.dropped_replicas

    @pytest.mark.parametrize(
        "setup_router",
        [
            {
                "enable_strict_max_ongoing_requests": True,
                "enable_queue_len_cache": True,
            },
        ],
        indirect=True,
    )
    async def test_replica_actor_unavailable(
        self, setup_router: Tuple[AsyncioRouter, FakeRequestRouter]
    ):
        router, fake_request_router = setup_router
        # Two replicas
        d_id = DeploymentID(name="test")
        r1_id = ReplicaID(unique_id="r1", deployment_id=d_id)
        r2_id = ReplicaID(unique_id="r2", deployment_id=d_id)

        # First request is sent to r1, cache should be populated with r1:5
        fake_request_router.set_replica_to_return(
            FakeReplica(
                r1_id,
                queue_len_info=ReplicaQueueLengthInfo(
                    accepted=True, num_ongoing_requests=5
                ),
            )
        )
        replica_result = await router.assign_request(dummy_request_metadata())
        assert replica_result._replica_id == r1_id
        # Cache should have R1:5
        assert fake_request_router.replica_queue_len_cache.get(r1_id) == 5
        assert fake_request_router.replica_queue_len_cache.get(r2_id) is None

        # Second request is sent to r2, cache should be populated with r2:10
        fake_request_router.set_replica_to_return(
            FakeReplica(
                r2_id,
                queue_len_info=ReplicaQueueLengthInfo(
                    accepted=True, num_ongoing_requests=10
                ),
            )
        )
        replica_result = await router.assign_request(dummy_request_metadata())
        assert replica_result._replica_id == r2_id
        # Cache should have R1:5, R2:10
        assert fake_request_router.replica_queue_len_cache.get(r1_id) == 5
        assert fake_request_router.replica_queue_len_cache.get(r2_id) == 10

        # Third request is sent to r1 again, but system message yields
        # an ActorUnavailableError
        fake_request_router.set_replica_to_return(
            FakeReplica(
                r1_id,
                error=ActorUnavailableError(error_message="unavailable", actor_id=None),
            )
        )
        fake_request_router.set_replica_to_return_on_retry(
            FakeReplica(
                r2_id,
                queue_len_info=ReplicaQueueLengthInfo(
                    accepted=True, num_ongoing_requests=15
                ),
            )
        )
        await router.assign_request(dummy_request_metadata())
        # R1 should be REMOVED from cache, cache should now be R2:15
        assert fake_request_router.replica_queue_len_cache.get(r1_id) is None
        assert fake_request_router.replica_queue_len_cache.get(r2_id) == 15

    @pytest.mark.parametrize(
        "setup_router",
        [
            {
                "enable_strict_max_ongoing_requests": True,
            },
        ],
        indirect=True,
    )
    @pytest.mark.parametrize("is_streaming", [False, True])
    async def test_on_request_routed(
        self,
        setup_router: Tuple[AsyncioRouter, FakeRequestRouter],
        is_streaming: bool,
    ):
        """Test that on_request_routed is called when a request is routed."""
        router, fake_request_router = setup_router

        # Before the request is routed, on_request_routed_called should be False.
        assert fake_request_router.on_request_routed_called is False

        r1_id = ReplicaID(
            unique_id="test-replica-1", deployment_id=DeploymentID(name="test")
        )
        replica = FakeReplica(
            r1_id,
            queue_len_info=ReplicaQueueLengthInfo(
                accepted=True, num_ongoing_requests=1
            ),
        )
        fake_request_router.set_replica_to_return(replica)

        replica_result = await router.assign_request(dummy_request_metadata())
        assert replica_result._replica_id == r1_id

        # After the request is routed, on_request_routed_called should be False.
        assert fake_request_router.on_request_routed_called is True


@pytest.mark.asyncio
class TestChooseReplica:
    """Tests for choose_replica() and dispatch() flow."""

    async def test_choose_replica_raises_deployment_unavailable(
        self, setup_router: Tuple[AsyncioRouter, FakeRequestRouter]
    ):
        """choose_replica fails when deployment is marked unavailable."""
        router, _ = setup_router
        router._deployment_available = False

        request_metadata = RequestMetadata(
            request_id="test-request-1",
            internal_request_id="test-internal-request-1",
        )

        with pytest.raises(DeploymentUnavailableError):
            async with router.choose_replica(request_metadata):
                pass

    async def test_choose_replica_raises_backpressure(
        self, setup_router: Tuple[AsyncioRouter, FakeRequestRouter]
    ):
        """choose_replica raises BackPressureError when queue limit is reached."""
        router, _ = setup_router
        router.update_deployment_config(DeploymentConfig(max_queued_requests=1))
        # Bump num_queued_requests to 1 to simulate a full queue
        router._metrics_manager.inc_num_queued_requests()

        request_metadata = RequestMetadata(
            request_id="test-request-1",
            internal_request_id="test-internal-request-1",
        )

        with pytest.raises(BackPressureError):
            async with router.choose_replica(request_metadata):
                pass

    async def test_choose_without_dispatch_releases_slot(
        self, setup_router: Tuple[AsyncioRouter, FakeRequestRouter]
    ):
        """Test that exiting context without dispatch releases the slot."""
        router, fake_request_router = setup_router

        r1_id = ReplicaID(
            unique_id="test-replica-1", deployment_id=DeploymentID(name="test")
        )
        replica = FakeReplica(r1_id)
        fake_request_router.set_replica_to_return(replica)

        request_metadata = RequestMetadata(
            request_id="test-request-1",
            internal_request_id="test-internal-request-1",
        )

        slot_token = None
        async with router.choose_replica(request_metadata) as selection:
            slot_token = selection._slot_token
            assert slot_token in replica._reserved_slots
            # Exit without calling dispatch

        # After context exit, slot should be released
        assert slot_token not in replica._reserved_slots

    async def test_choose_without_dispatch_does_not_fire_on_request_completed(
        self, setup_router: Tuple[AsyncioRouter, FakeRequestRouter]
    ):
        """on_request_completed shouldn't fire if the caller exits without dispatching."""
        router, fake_request_router = setup_router

        r1_id = ReplicaID(
            unique_id="test-replica-1", deployment_id=DeploymentID(name="test")
        )
        fake_request_router.set_replica_to_return(FakeReplica(r1_id))

        async with router.choose_replica(dummy_request_metadata()):
            pass  # Exit without dispatch.

        assert fake_request_router.on_request_routed_called is False
        assert fake_request_router.completed_requests == []

    async def test_choose_with_exception_releases_slot(
        self, setup_router: Tuple[AsyncioRouter, FakeRequestRouter]
    ):
        """Test that exception in context releases the slot."""
        router, fake_request_router = setup_router

        r1_id = ReplicaID(
            unique_id="test-replica-1", deployment_id=DeploymentID(name="test")
        )
        replica = FakeReplica(r1_id)
        fake_request_router.set_replica_to_return(replica)

        request_metadata = RequestMetadata(
            request_id="test-request-1",
            internal_request_id="test-internal-request-1",
        )

        slot_token = None
        with pytest.raises(RuntimeError):
            async with router.choose_replica(request_metadata) as selection:
                slot_token = selection._slot_token
                assert slot_token in replica._reserved_slots
                raise RuntimeError("Test exception")

        # After exception, slot should be released
        assert slot_token not in replica._reserved_slots

    @pytest.mark.parametrize("is_streaming", [False, True])
    async def test_choose_and_dispatch(
        self,
        setup_router: Tuple[AsyncioRouter, FakeRequestRouter],
        is_streaming: bool,
    ):
        """Happy path: selection fields are populated, dispatch sends with
        with_rejection=False (slot reservation replaces the rejection
        round-trip), and the slot is released on exit."""
        router, fake_request_router = setup_router

        r1_id = ReplicaID(
            unique_id="test-replica-1", deployment_id=DeploymentID(name="test")
        )
        replica = FakeReplica(r1_id)
        fake_request_router.set_replica_to_return(replica)

        request_metadata = RequestMetadata(
            request_id="test-request-1",
            internal_request_id="test-internal-request-1",
            call_method="test_method",
            is_streaming=is_streaming,
        )

        async with router.choose_replica(request_metadata) as selection:
            assert selection.replica_id == r1_id.unique_id
            assert selection._replica == replica
            assert selection._method_name == "test_method"
            assert selection._slot_token in replica._reserved_slots

            replica_result = await router.dispatch(selection, request_metadata)
            assert replica_result._replica_id == r1_id
            assert replica_result._is_generator_object == is_streaming

        assert selection._slot_token not in replica._reserved_slots
        assert len(replica._requests_sent) == 1
        assert replica._requests_sent[0]["with_rejection"] is False

    async def test_multiple_sequential_selections(
        self, setup_router: Tuple[AsyncioRouter, FakeRequestRouter]
    ):
        """Test multiple sequential choose_replica() and dispatch() calls."""
        router, fake_request_router = setup_router

        r1_id = ReplicaID(
            unique_id="test-replica-1", deployment_id=DeploymentID(name="test")
        )
        replica = FakeReplica(r1_id)
        fake_request_router.set_replica_to_return(replica)

        for i in range(3):
            request_metadata = RequestMetadata(
                request_id=f"test-request-{i}",
                internal_request_id=f"test-internal-request-{i}",
            )

            async with router.choose_replica(request_metadata) as selection:
                replica_result = await router.dispatch(selection, request_metadata)
                assert replica_result._replica_id == r1_id

        # All slots should have been created and used
        assert replica._slot_counter == 3

    @pytest.mark.parametrize(
        "setup_router",
        [{"enable_queue_len_cache": True}],
        indirect=True,
    )
    async def test_cache_updated_on_choose_replica(
        self, setup_router: Tuple[AsyncioRouter, FakeRequestRouter]
    ):
        """Test that queue length cache is updated when choosing a replica."""
        router, fake_request_router = setup_router

        r1_id = ReplicaID(
            unique_id="test-replica-1", deployment_id=DeploymentID(name="test")
        )
        replica = FakeReplica(r1_id)
        fake_request_router.set_replica_to_return(replica)

        request_metadata = RequestMetadata(
            request_id="test-request-1",
            internal_request_id="test-internal-request-1",
        )

        # Initially cache should be empty
        assert fake_request_router.replica_queue_len_cache.get(r1_id) is None

        # Choose replica should update cache to 1
        async with router.choose_replica(request_metadata) as selection:
            assert fake_request_router.replica_queue_len_cache.get(r1_id) == 1

            # Dispatch should NOT increment again (already counted in choose)
            await router.dispatch(selection, request_metadata)
            assert fake_request_router.replica_queue_len_cache.get(r1_id) == 1

        # After dispatch, cache remains incremented while the request is in flight.
        # It is decremented when request completion is observed.
        assert fake_request_router.replica_queue_len_cache.get(r1_id) == 1

    @pytest.mark.parametrize(
        "setup_router",
        [{"enable_queue_len_cache": True}],
        indirect=True,
    )
    async def test_current_loop_dispatch_marks_selection_before_task_runs(
        self, setup_router: Tuple[AsyncioRouter, FakeRequestRouter]
    ):
        """Wrapper dispatch should consume the selection before its task runs."""
        router, fake_request_router = setup_router

        r1_id = ReplicaID(
            unique_id="test-replica-1", deployment_id=DeploymentID(name="test")
        )
        replica = FakeReplica(r1_id)
        fake_request_router.set_replica_to_return(replica)

        current_loop_router = CurrentLoopRouter.__new__(CurrentLoopRouter)
        current_loop_router._asyncio_loop = asyncio.get_running_loop()
        current_loop_router._asyncio_router = router

        request_metadata = RequestMetadata(
            request_id="test-request-1",
            internal_request_id="test-internal-request-1",
        )

        async with router.choose_replica(request_metadata) as selection:
            dispatch_task = current_loop_router.dispatch(selection, request_metadata)
            assert selection._dispatched
            assert fake_request_router.replica_queue_len_cache.get(r1_id) == 1

        assert fake_request_router.replica_queue_len_cache.get(r1_id) == 1
        assert len(replica._requests_sent) == 0

        replica_result = await dispatch_task
        assert fake_request_router.replica_queue_len_cache.get(r1_id) == 1

        replica_result.fire_done_callbacks()
        await asyncio.sleep(0)
        assert fake_request_router.replica_queue_len_cache.get(r1_id) == 0

    @pytest.mark.parametrize(
        "setup_router",
        [{"enable_queue_len_cache": True}],
        indirect=True,
    )
    async def test_current_loop_dispatch_failure_releases_cache(
        self, setup_router: Tuple[AsyncioRouter, FakeRequestRouter]
    ):
        """If wrapper dispatch fails after consuming a selection, it owns cleanup."""
        router, fake_request_router = setup_router

        r1_id = ReplicaID(
            unique_id="test-replica-1", deployment_id=DeploymentID(name="test")
        )
        replica = FakeReplica(r1_id)
        fake_request_router.set_replica_to_return(replica)

        current_loop_router = CurrentLoopRouter.__new__(CurrentLoopRouter)
        current_loop_router._asyncio_loop = asyncio.get_running_loop()
        current_loop_router._asyncio_router = router

        request_metadata = RequestMetadata(
            request_id="test-request-1",
            internal_request_id="test-internal-request-1",
        )

        async with router.choose_replica(request_metadata) as selection:
            fake_request_router._replica_to_return = None
            dispatch_task = current_loop_router.dispatch(selection, request_metadata)
            assert selection._dispatched

        assert fake_request_router.replica_queue_len_cache.get(r1_id) == 1

        with pytest.raises(ReplicaUnavailableError):
            await dispatch_task
        assert fake_request_router.replica_queue_len_cache.get(r1_id) == 0

    @pytest.mark.parametrize(
        "setup_router",
        [{"enable_queue_len_cache": True}],
        indirect=True,
    )
    async def test_cache_decremented_on_choose_without_dispatch(
        self, setup_router: Tuple[AsyncioRouter, FakeRequestRouter]
    ):
        """Test that cache is decremented when choose exits without dispatch."""
        router, fake_request_router = setup_router

        r1_id = ReplicaID(
            unique_id="test-replica-1", deployment_id=DeploymentID(name="test")
        )
        replica = FakeReplica(r1_id)
        fake_request_router.set_replica_to_return(replica)

        request_metadata = RequestMetadata(
            request_id="test-request-1",
            internal_request_id="test-internal-request-1",
        )

        # Choose replica without dispatch
        async with router.choose_replica(request_metadata):
            # Cache should be 1 (reservation)
            assert fake_request_router.replica_queue_len_cache.get(r1_id) == 1
            # Exit without calling dispatch

        # After context exit without dispatch, cache should be decremented to 0
        assert fake_request_router.replica_queue_len_cache.get(r1_id) == 0

    @pytest.mark.parametrize(
        "setup_router",
        [{"enable_queue_len_cache": True}],
        indirect=True,
    )
    async def test_choose_replica_retries_when_reservation_rejected(
        self, setup_router: Tuple[AsyncioRouter, FakeRequestRouter]
    ):
        """choose_replica should only yield after replica-side capacity is reserved."""
        router, fake_request_router = setup_router

        r1_id = ReplicaID(
            unique_id="test-replica-1", deployment_id=DeploymentID(name="test")
        )
        r2_id = ReplicaID(
            unique_id="test-replica-2", deployment_id=DeploymentID(name="test")
        )
        r1 = FakeReplica(r1_id)
        r2 = FakeReplica(r2_id)
        r1._reject_reservation = True
        fake_request_router.set_replica_to_return(r1)
        fake_request_router.set_replica_to_return_on_retry(r2)

        request_metadata = RequestMetadata(
            request_id="test-request-1",
            internal_request_id="test-internal-request-1",
        )

        async with router.choose_replica(request_metadata) as selection:
            assert selection._replica == r2
            assert fake_request_router.replica_queue_len_cache.get(r1_id) == 1
            assert fake_request_router.replica_queue_len_cache.get(r2_id) == 1

        assert fake_request_router.replica_queue_len_cache.get(r2_id) == 0

    @pytest.mark.parametrize(
        "setup_router",
        [{"enable_queue_len_cache": True}],
        indirect=True,
    )
    async def test_concurrent_choose_replica_updates_cache(
        self, setup_router: Tuple[AsyncioRouter, FakeRequestRouter]
    ):
        """Test that concurrent choose_replica calls correctly update cache."""
        router, fake_request_router = setup_router

        r1_id = ReplicaID(
            unique_id="test-replica-1", deployment_id=DeploymentID(name="test")
        )
        replica = FakeReplica(r1_id)
        fake_request_router.set_replica_to_return(replica)

        # Start 3 concurrent choose_replica operations
        async def choose_and_hold(request_id: str):
            metadata = RequestMetadata(
                request_id=request_id,
                internal_request_id=request_id,
            )
            async with router.choose_replica(metadata) as selection:
                # Hold the selection
                await asyncio.sleep(0.01)
                return selection

        # Create tasks
        tasks = [asyncio.create_task(choose_and_hold(f"request-{i}")) for i in range(3)]

        # Wait a bit for all to enter context
        await asyncio.sleep(0.005)

        # Cache should reflect all 3 reservations
        assert fake_request_router.replica_queue_len_cache.get(r1_id) == 3

        # Let them all exit
        await asyncio.gather(*tasks)

        # After all exit without dispatch, cache should be 0
        assert fake_request_router.replica_queue_len_cache.get(r1_id) == 0

    async def test_reserved_slots_gauge_increments_and_decrements(
        self, setup_router: Tuple[AsyncioRouter, FakeRequestRouter]
    ):
        """Gauge goes 0 → 1 → 0 across choose+dispatch and choose-without-dispatch.

        Asserts the underlying `.set()` calls too, so bumping the counter
        without pushing to the gauge is caught.
        """
        router, fake_request_router = setup_router

        gauge_mock = Mock()
        router._metrics_manager._reserved_slots_gauge = gauge_mock

        r1_id = ReplicaID(
            unique_id="test-replica-1", deployment_id=DeploymentID(name="test")
        )
        replica = FakeReplica(r1_id)
        fake_request_router.set_replica_to_return(replica)

        # Choose + dispatch path.
        request_metadata = RequestMetadata(
            request_id="test-request-1",
            internal_request_id="test-internal-request-1",
        )
        assert router._metrics_manager._num_reserved_slots == 0
        async with router.choose_replica(request_metadata) as selection:
            assert router._metrics_manager._num_reserved_slots == 1
            gauge_mock.set.assert_called_with(1)
            await router.dispatch(selection, request_metadata)

        assert router._metrics_manager._num_reserved_slots == 0
        gauge_mock.set.assert_called_with(0)

        # Choose without dispatch: same gauge transition must hold.
        gauge_mock.reset_mock()
        request_metadata = RequestMetadata(
            request_id="test-request-2",
            internal_request_id="test-internal-request-2",
        )
        async with router.choose_replica(request_metadata):
            assert router._metrics_manager._num_reserved_slots == 1
            gauge_mock.set.assert_called_with(1)

        assert router._metrics_manager._num_reserved_slots == 0
        gauge_mock.set.assert_called_with(0)

    async def test_release_failure_does_not_leak_reserved_slots_metric(
        self, setup_router: Tuple[AsyncioRouter, FakeRequestRouter]
    ):
        """If the release call raises during choose_replica cleanup, the
        reserved-slots gauge must still decrement — otherwise it leaks for the
        process lifetime — and the exception must not propagate to the caller.
        """
        router, fake_request_router = setup_router

        r1_id = ReplicaID(
            unique_id="test-replica-1", deployment_id=DeploymentID(name="test")
        )
        replica = FakeReplica(r1_id)
        replica.set_release_slot_error(RuntimeError("simulated release failure"))
        fake_request_router.set_replica_to_return(replica)

        request_metadata = RequestMetadata(
            request_id="test-request-1",
            internal_request_id="test-internal-request-1",
        )

        before = router._metrics_manager._num_reserved_slots
        async with router.choose_replica(request_metadata):
            assert router._metrics_manager._num_reserved_slots == before + 1
        # Exit ran cleanly even though release raised.
        assert router._metrics_manager._num_reserved_slots == before

    async def test_inc_reserved_slots_failure_releases_slot(
        self, setup_router: Tuple[AsyncioRouter, FakeRequestRouter]
    ):
        """If gauge.set() throws during inc, cleanup must still release the
        slot and rebalance the counter."""
        router, fake_request_router = setup_router

        r1_id = ReplicaID(
            unique_id="test-replica-1", deployment_id=DeploymentID(name="test")
        )
        replica = FakeReplica(r1_id)
        fake_request_router.set_replica_to_return(replica)

        # Raise only on inc's set(1); the cleanup's set(0) must still run.
        gauge_mock = Mock()
        gauge_mock.set = Mock(
            side_effect=lambda v: (_ for _ in ()).throw(RuntimeError("gauge boom"))
            if v > 0
            else None
        )
        router._metrics_manager._reserved_slots_gauge = gauge_mock

        with pytest.raises(RuntimeError, match="gauge boom"):
            async with router.choose_replica(dummy_request_metadata()):
                pass

        assert replica._reserved_slots == set()
        assert router._metrics_manager._num_reserved_slots == 0

    @pytest.mark.parametrize(
        "callback_result,expect_release",
        [
            pytest.param(None, False, id="actor_success"),
            pytest.param(TaskCancelledError(), True, id="actor_cancellation"),
            pytest.param(
                Mock(cancelled=Mock(return_value=False)),
                False,
                id="grpc_success",
            ),
            pytest.param(
                Mock(cancelled=Mock(return_value=True)),
                True,
                id="grpc_cancellation",
            ),
        ],
    )
    async def test_dispatch_release_slot_fires_only_on_cancellation(
        self,
        setup_router: Tuple[AsyncioRouter, FakeRequestRouter],
        callback_result,
        expect_release,
    ):
        """Done-callback fires release_slot only when the dispatched
        request was cancelled. On success the replica already consumed the
        slot when try_send_request arrived, so a follow-up RPC would be a
        no-op per dispatched request. Covers both transports."""
        router, fake_request_router = setup_router

        r1_id = ReplicaID(
            unique_id="test-replica-1", deployment_id=DeploymentID(name="test")
        )
        replica = FakeReplica(r1_id)
        fake_request_router.set_replica_to_return(replica)

        release_slot_calls: List[str] = []
        original_release_slot = replica.release_slot

        async def spy_release_slot(slot_token):
            release_slot_calls.append(slot_token)
            return await original_release_slot(slot_token)

        replica.release_slot = spy_release_slot

        request_metadata = dummy_request_metadata()
        async with router.choose_replica(request_metadata) as selection:
            slot_token = selection._slot_token
            result = await router.dispatch(selection, request_metadata)

        # try_send_request consumed the slot synchronously; the done-callback
        # hasn't fired yet, so no release_slot call should have happened.
        assert release_slot_calls == []

        result.fire_done_callbacks(result=callback_result)
        # Drain the call_soon_threadsafe → create_task chain.
        for _ in range(2):
            await asyncio.sleep(0)

        assert release_slot_calls == ([slot_token] if expect_release else [])

    async def test_dispatch_failure_is_not_masked_by_release_failure(
        self, setup_router: Tuple[AsyncioRouter, FakeRequestRouter]
    ):
        """If both `try_send_request` and the cleanup `release_slot` raise, the
        caller must see the original dispatch error, not the release error.
        """
        router, fake_request_router = setup_router

        r1_id = ReplicaID(
            unique_id="test-replica-1", deployment_id=DeploymentID(name="test")
        )
        dispatch_error = RuntimeError("dispatch boom")
        replica = FakeReplica(r1_id, error=dispatch_error)
        replica.set_release_slot_error(RuntimeError("release boom"))
        fake_request_router.set_replica_to_return(replica)

        request_metadata = RequestMetadata(
            request_id="test-request-1",
            internal_request_id="test-internal-request-1",
        )

        async with router.choose_replica(request_metadata) as selection:
            with pytest.raises(RuntimeError, match="dispatch boom"):
                await router.dispatch(selection, request_metadata)

    async def test_dispatch_replica_unavailable(
        self, setup_router: Tuple[AsyncioRouter, FakeRequestRouter]
    ):
        """Test dispatch() raises error when replica becomes unavailable."""
        router, fake_request_router = setup_router

        r1_id = ReplicaID(
            unique_id="test-replica-1", deployment_id=DeploymentID(name="test")
        )
        replica = FakeReplica(r1_id)
        fake_request_router.set_replica_to_return(replica)

        request_metadata = RequestMetadata(
            request_id="test-request-1",
            internal_request_id="test-internal-request-1",
        )

        async with router.choose_replica(request_metadata) as selection:
            # Simulate replica becoming unavailable
            fake_request_router._replica_to_return = None

            # dispatch should raise ReplicaUnavailableError
            with pytest.raises(ReplicaUnavailableError):
                await router.dispatch(selection, request_metadata)

    async def test_multiple_dispatch_calls_fail(
        self, setup_router: Tuple[AsyncioRouter, FakeRequestRouter]
    ):
        """A ReplicaSelection can only be dispatched once."""
        router, fake_request_router = setup_router

        r1_id = ReplicaID(
            unique_id="test-replica-1", deployment_id=DeploymentID(name="test")
        )
        replica = FakeReplica(r1_id)
        fake_request_router.set_replica_to_return(replica)

        request_metadata = RequestMetadata(
            request_id="test-request-1",
            internal_request_id="test-internal-request-1",
        )

        async with router.choose_replica(request_metadata) as selection:
            await router.dispatch(selection, request_metadata)

            with pytest.raises(RuntimeError, match="already been dispatched"):
                await router.dispatch(selection, request_metadata)

        assert len(replica._requests_sent) == 1

    async def test_choose_replica_no_reserve_skips_slot_reservation(
        self, setup_router: Tuple[AsyncioRouter, FakeRequestRouter]
    ):
        """``_reserve=False`` yields a selection with no slot token, bypasses
        ``reserve_slot`` on the replica, and ``dispatch`` on that selection
        raises since there's no reservation to consume."""
        router, fake_request_router = setup_router

        r1_id = ReplicaID(
            unique_id="test-replica-1", deployment_id=DeploymentID(name="test")
        )
        replica = FakeReplica(r1_id)
        fake_request_router._replicas_list = [replica]

        async def fake_choose_replicas(candidate_replicas, pending_request=None):
            return [candidate_replicas]

        fake_request_router.choose_replicas = fake_choose_replicas

        request_metadata = dummy_request_metadata()

        async with router.choose_replica(request_metadata, _reserve=False) as selection:
            assert selection._replica is replica
            assert selection._slot_token is None
            assert replica._reserved_slots == set()

            with pytest.raises(RuntimeError, match="_reserve=False"):
                await router.dispatch(selection, request_metadata)

        assert replica._reserved_slots == set()
        assert replica._slot_counter == 0


def running_replica_info(replica_id: ReplicaID) -> RunningReplicaInfo:
    return RunningReplicaInfo(
        replica_id=replica_id,
        node_id="node_id",
        node_ip="node_ip",
        availability_zone="some-az",
        actor_name=replica_id.to_full_id_str(),
        max_ongoing_requests=1,
    )


class TestRouterMetricsManager:
    @pytest.mark.asyncio
    async def test_num_router_requests(self):
        tags = {
            "deployment": "a",
            "application": "b",
            "route": "/alice",
            "handle": "random_handle",
            "actor_id": "random_actor",
        }
        metrics_manager = RouterMetricsManager(
            DeploymentID(name="a", app_name="b"),
            "random_handle",
            "random_actor",
            DeploymentHandleSource.UNKNOWN,
            Mock(),
            FakeCounter(
                tag_keys=("deployment", "route", "application", "handle", "actor_id")
            ),
            FakeGauge(tag_keys=("deployment", "application", "handle", "actor_id")),
            FakeGauge(tag_keys=("deployment", "application", "handle", "actor_id")),
            event_loop=asyncio.get_event_loop(),
        )
        assert metrics_manager.num_router_requests.get_count(tags) is None

        n = random.randint(1, 10)
        for _ in range(n):
            metrics_manager.inc_num_total_requests(route="/alice")

        await asyncio.sleep(RAY_SERVE_METRICS_EXPORT_INTERVAL_MS * 2 / 1000)
        assert metrics_manager.num_router_requests.get_count(tags) == n

    @pytest.mark.asyncio
    async def test_num_queued_requests_gauge(self):
        tags = {
            "deployment": "a",
            "application": "b",
            "handle": "random_handle",
            "actor_id": "random_actor",
        }
        metrics_manager = RouterMetricsManager(
            DeploymentID(name="a", app_name="b"),
            "random_handle",
            "random_actor",
            DeploymentHandleSource.UNKNOWN,
            Mock(),
            FakeCounter(
                tag_keys=("deployment", "route", "application", "handle", "actor_id")
            ),
            FakeGauge(tag_keys=("deployment", "application", "handle", "actor_id")),
            FakeGauge(tag_keys=("deployment", "application", "handle", "actor_id")),
            event_loop=asyncio.get_event_loop(),
        )
        assert metrics_manager.num_queued_requests_gauge.get_value(tags) == 0

        n, m = random.randint(0, 10), random.randint(0, 5)
        for _ in range(n):
            metrics_manager.inc_num_queued_requests()
        await asyncio.sleep(RAY_SERVE_METRICS_EXPORT_INTERVAL_MS * 2 / 1000)
        assert metrics_manager.num_queued_requests_gauge.get_value(tags) == n
        for _ in range(m):
            metrics_manager.dec_num_queued_requests()

        await asyncio.sleep(RAY_SERVE_METRICS_EXPORT_INTERVAL_MS * 2 / 1000)
        assert metrics_manager.num_queued_requests_gauge.get_value(tags) == n - m

    @pytest.mark.asyncio
    async def test_track_requests_sent_to_replicas(self):
        d_id = DeploymentID(name="a", app_name="b")
        metrics_manager = RouterMetricsManager(
            d_id,
            "random",
            "random_actor",
            DeploymentHandleSource.UNKNOWN,
            Mock(),
            FakeCounter(
                tag_keys=("deployment", "route", "application", "handle", "actor_id")
            ),
            FakeGauge(tag_keys=("deployment", "application", "handle", "actor_id")),
            FakeGauge(tag_keys=("deployment", "application", "handle", "actor_id")),
            event_loop=asyncio.get_event_loop(),
        )

        # r1: number requests -> 0, removed from list of running replicas -> prune
        # r2: number requests -> 0, remains on list of running replicas -> don't prune
        # r3: number requests > 0, removed from list of running replicas -> don't prune
        # r4: number requests > 0, remains on list of running replicas -> don't prune
        replica_ids = [
            ReplicaID(unique_id=f"test-replica-{i}", deployment_id=d_id)
            for i in range(1, 5)
        ]
        r1, r2, r3, r4 = replica_ids

        # ri has i requests
        for i in range(4):
            for _ in range(i + 1):
                metrics_manager.inc_num_running_requests_for_replica(replica_ids[i])
        await asyncio.sleep(RAY_SERVE_METRICS_EXPORT_INTERVAL_MS * 2 / 1000)

        # All 4 replicas should have a positive number of requests
        for i, r in enumerate(replica_ids):
            assert metrics_manager.num_requests_sent_to_replicas[r] == i + 1
        assert (
            metrics_manager.num_running_requests_gauge.get_value(
                {
                    "deployment": "a",
                    "application": "b",
                    "handle": "random",
                    "actor_id": "random_actor",
                }
            )
            == 10
        )

        # Requests at r1 and r2 drop to 0
        for _ in range(1):
            metrics_manager.dec_num_running_requests_for_replica(r1)
        for _ in range(2):
            metrics_manager.dec_num_running_requests_for_replica(r2)
        await asyncio.sleep(RAY_SERVE_METRICS_EXPORT_INTERVAL_MS * 2 / 1000)
        assert metrics_manager.num_requests_sent_to_replicas[r1] == 0
        assert metrics_manager.num_requests_sent_to_replicas[r2] == 0

        # 3 requests finished processing
        assert (
            metrics_manager.num_running_requests_gauge.get_value(
                {
                    "deployment": "a",
                    "application": "b",
                    "handle": "random",
                    "actor_id": "random_actor",
                }
            )
            == 7
        )

        # Running replicas reduces to [r2, r4]
        metrics_manager._update_running_replicas(
            [
                running_replica_info(r2),
                running_replica_info(r4),
            ]
        )
        await asyncio.sleep(RAY_SERVE_METRICS_EXPORT_INTERVAL_MS * 2 / 1000)

        # Only r1 should be pruned, the rest should still be tracked.
        assert r1 not in metrics_manager.num_requests_sent_to_replicas
        assert r2 in metrics_manager.num_requests_sent_to_replicas
        assert r3 in metrics_manager.num_requests_sent_to_replicas
        assert r4 in metrics_manager.num_requests_sent_to_replicas

    @pytest.mark.asyncio
    async def test_should_send_scaled_to_zero_optimized_push(self):
        metrics_manager = RouterMetricsManager(
            DeploymentID(name="a", app_name="b"),
            "random",
            "random_actor",
            DeploymentHandleSource.UNKNOWN,
            Mock(),
            FakeCounter(
                tag_keys=("deployment", "route", "application", "handle", "actor_id")
            ),
            FakeGauge(tag_keys=("deployment", "application", "handle", "actor_id")),
            FakeGauge(tag_keys=("deployment", "application", "handle", "actor_id")),
            event_loop=asyncio.get_event_loop(),
        )

        # Not an autoscaling deployment, should not push metrics
        assert not metrics_manager.should_send_scaled_to_zero_optimized_push(0)

        # No queued requests at the handle, should not push metrics
        metrics_manager._deployment_config = DeploymentConfig(
            autoscaling_config=AutoscalingConfig()
        )
        assert not metrics_manager.should_send_scaled_to_zero_optimized_push(0)

        # Current number of replicas is non-zero, should not push metrics
        metrics_manager.inc_num_queued_requests()
        assert not metrics_manager.should_send_scaled_to_zero_optimized_push(1)

        # All 3 conditions satisfied, should push metrics
        assert metrics_manager.should_send_scaled_to_zero_optimized_push(0)

    @pytest.mark.asyncio
    @patch(
        "ray.serve._private.router.RAY_SERVE_COLLECT_AUTOSCALING_METRICS_ON_HANDLE", "1"
    )
    async def test_push_autoscaling_metrics_to_controller(self):
        timer = MockTimer()
        start = random.randint(50, 100)
        timer.reset(start)
        deployment_id = DeploymentID(name="a", app_name="b")
        handle_id = "random"
        self_actor_id = "abc"
        mock_controller_handle = Mock()

        with patch("time.time", new=timer.time):
            metrics_manager = RouterMetricsManager(
                deployment_id,
                handle_id,
                self_actor_id,
                DeploymentHandleSource.PROXY,
                mock_controller_handle,
                FakeCounter(
                    tag_keys=(
                        "deployment",
                        "route",
                        "application",
                        "handle",
                        "actor_id",
                    )
                ),
                FakeGauge(tag_keys=("deployment", "application", "handle", "actor_id")),
                FakeGauge(tag_keys=("deployment", "application", "handle", "actor_id")),
                event_loop=asyncio.get_event_loop(),
            )
            metrics_manager._deployment_config = DeploymentConfig(
                autoscaling_config=AutoscalingConfig()
            )

            # Set up some requests
            n = random.randint(0, 5)
            replica_ids = [
                ReplicaID(get_random_string(), DeploymentID("d", "a")) for _ in range(3)
            ]
            running_requests = defaultdict(int)
            for _ in range(n):
                metrics_manager.inc_num_queued_requests()
            for _ in range(20):
                r = random.choice(replica_ids)
                running_requests[r] += 1
                metrics_manager.inc_num_running_requests_for_replica(r)

            # Check metrics are pushed correctly (compressed)
            metrics_manager.push_autoscaling_metrics_to_controller()
            mock_controller_handle.record_autoscaling_metrics_from_handle.remote.assert_called_once()
            (
                compressed,
            ) = mock_controller_handle.record_autoscaling_metrics_from_handle.remote.call_args[
                0
            ]
            assert isinstance(compressed, bytes)
            handle_metric_report = decompress_metric_report(compressed)
            assert handle_metric_report.deployment_id == deployment_id
            assert handle_metric_report.handle_id == handle_id

    @pytest.mark.skipif(
        not RAY_SERVE_COLLECT_AUTOSCALING_METRICS_ON_HANDLE,
        reason="Tests handle metrics behavior.",
    )
    @pytest.mark.asyncio
    @patch(
        "ray.serve._private.router.RAY_SERVE_AUTOSCALING_METRIC_RECORD_INTERVAL_FACTOR",
        0.001,
    )
    async def test_memory_cleared(self):
        deployment_id = DeploymentID(name="a", app_name="b")
        metrics_manager = RouterMetricsManager(
            deployment_id,
            "some_handle",
            "some_actor",
            DeploymentHandleSource.PROXY,
            Mock(),
            FakeCounter(
                tag_keys=(
                    "deployment",
                    "route",
                    "application",
                    "handle",
                    "actor_id",
                )
            ),
            FakeGauge(tag_keys=("deployment", "application", "handle", "actor_id")),
            FakeGauge(tag_keys=("deployment", "application", "handle", "actor_id")),
            event_loop=asyncio.get_event_loop(),
        )
        metrics_manager.update_deployment_config(
            deployment_config=DeploymentConfig(
                autoscaling_config=AutoscalingConfig(
                    metrics_interval_s=0.005, look_back_period_s=0.01
                )
            ),
            curr_num_replicas=0,
        )

        r1 = ReplicaID("r1", deployment_id)
        r2 = ReplicaID("r2", deployment_id)
        r3 = ReplicaID("r3", deployment_id)

        def check_database(expected: Set[ReplicaID]):
            assert set(metrics_manager.metrics_store.data) == expected
            return True

        # r1: 1
        metrics_manager.inc_num_running_requests_for_replica(r1)
        await async_wait_for_condition(
            check_database, expected={r1, QUEUED_REQUESTS_KEY}
        )

        # r1: 1, r2: 0
        metrics_manager.inc_num_running_requests_for_replica(r2)
        await async_wait_for_condition(
            check_database, expected={r1, r2, QUEUED_REQUESTS_KEY}
        )
        metrics_manager.dec_num_running_requests_for_replica(r2)

        # r1: 1, r2: 0, r3: 0
        metrics_manager.inc_num_running_requests_for_replica(r3)
        await async_wait_for_condition(
            check_database, expected={r1, r2, r3, QUEUED_REQUESTS_KEY}
        )
        metrics_manager.dec_num_running_requests_for_replica(r3)

        # update running replicas {r2}
        metrics_manager._update_running_replicas([running_replica_info(r2)])
        await async_wait_for_condition(
            check_database, expected={r1, r2, QUEUED_REQUESTS_KEY}
        )

    @pytest.mark.asyncio
    @patch(
        "ray.serve._private.router.RAY_SERVE_COLLECT_AUTOSCALING_METRICS_ON_HANDLE", "1"
    )
    @patch("ray.serve._private.router.MetricsPusher")
    async def test_update_deployment_config(self, metrics_pusher_mock):
        metrics_manager = RouterMetricsManager(
            DeploymentID(name="a", app_name="b"),
            "random",
            "random_actor",
            DeploymentHandleSource.UNKNOWN,
            Mock(),
            FakeCounter(
                tag_keys=("deployment", "route", "application", "handle", "actor_id")
            ),
            FakeGauge(tag_keys=("deployment", "application", "handle", "actor_id")),
            FakeGauge(tag_keys=("deployment", "application", "handle", "actor_id")),
            event_loop=asyncio.get_event_loop(),
        )

        # Without autoscaling config, do nothing
        metrics_manager.update_deployment_config(DeploymentConfig(), 0)
        metrics_manager.metrics_pusher.register_or_update_task.assert_not_called()

        # With autoscaling config, register or update task should be called
        metrics_manager.update_deployment_config(
            DeploymentConfig(autoscaling_config=AutoscalingConfig()), 0
        )
        metrics_manager.metrics_pusher.register_or_update_task.assert_called()


class TestSingletonThreadRouter:
    @pytest.fixture
    def setup_singleton_thread_router(
        self, setup_router: Tuple[AsyncioRouter, FakeRequestRouter]
    ) -> SingletonThreadRouter:
        asyncio_router, fake_request_router = setup_router

        # Mock ray.get_runtime_context() to avoid triggering Ray initialization
        with patch("ray.get_runtime_context") as mock_context:
            mock_context.return_value.get_actor_id.return_value = "test-actor-id"

            router = SingletonThreadRouter(
                controller_handle=Mock(),
                deployment_id=DeploymentID(name="test", app_name="test"),
                handle_id="test",
                self_actor_id="test",
                handle_source="test",
                request_router=fake_request_router,
                enable_strict_max_ongoing_requests=False,
                resolve_request_arg_func=Mock(),
                node_id="test-node-id",
                availability_zone="test-az",
                prefer_local_node_routing=False,
            )
        router._asyncio_router = asyncio_router
        return router

    def test_request_assignment(
        self,
        setup_router: Tuple[AsyncioRouter, FakeRequestRouter],
        setup_singleton_thread_router: SingletonThreadRouter,
    ):
        _, fake_request_router = setup_router
        thread_router = setup_singleton_thread_router

        r1_id = ReplicaID(
            unique_id="test-replica-1", deployment_id=DeploymentID(name="test")
        )
        replica = FakeReplica(r1_id)
        fake_request_router.set_replica_to_return(replica)

        request_metadata = RequestMetadata(
            request_id="test-request-1",
            internal_request_id="test-internal-request-1",
        )

        future = thread_router.assign_request(request_metadata)
        assert isinstance(future, concurrent.futures.Future)
        assert future.result()._replica_id == r1_id

    @pytest.mark.asyncio
    @pytest.mark.parametrize(
        "setup_router",
        [{"enable_queue_len_cache": True}],
        indirect=True,
    )
    async def test_dispatch_marks_selection_before_scheduled_coroutine_runs(
        self,
        setup_router: Tuple[AsyncioRouter, FakeRequestRouter],
        setup_singleton_thread_router: SingletonThreadRouter,
        monkeypatch,
    ):
        _, fake_request_router = setup_router
        thread_router = setup_singleton_thread_router

        r1_id = ReplicaID(
            unique_id="test-replica-1", deployment_id=DeploymentID(name="test")
        )
        replica = FakeReplica(r1_id)
        fake_request_router.set_replica_to_return(replica)

        pending_coros = []

        def delay_asyncio_call(coro):
            pending_coros.append(coro)
            return concurrent.futures.Future()

        monkeypatch.setattr(
            thread_router, "_wrap_asyncio_call_in_future", delay_asyncio_call
        )

        request_metadata = RequestMetadata(
            request_id="test-request-1",
            internal_request_id="test-internal-request-1",
        )

        async with thread_router.choose_replica(request_metadata) as selection:
            dispatch_future = thread_router.dispatch(selection, request_metadata)
            assert selection._dispatched

        assert fake_request_router.replica_queue_len_cache.get(r1_id) == 1
        assert len(pending_coros) == 1

        dispatch_future.cancel()
        pending_coros[0].close()

    @pytest.mark.asyncio
    async def test_choose_replica_cancel_during_entry_releases_slot(
        self,
        setup_router: Tuple[AsyncioRouter, FakeRequestRouter],
        setup_singleton_thread_router: SingletonThreadRouter,
        monkeypatch,
    ):
        """Cancelling the caller after __aenter__ reserved a slot on the router
        loop, but before the bridge takes ownership of the selection, must
        still release the slot via __aexit__.
        """
        fake_router, _ = setup_router
        thread_router = setup_singleton_thread_router
        outer_loop = asyncio.get_running_loop()

        fake_selection = ReplicaSelection(
            replica_id="fake-replica",
            node_ip="127.0.0.1",
            port=None,
            node_id="fake-node",
            availability_zone=None,
            replica_metadata={},
            _replica=None,
            _deployment_id=None,
            _request_metadata=None,
            _method_name="",
            _slot_token="",
        )

        reserved_slots = 0
        reserved_lock = threading.Lock()
        release_path: List[str] = []
        released = threading.Event()
        outer_task_holder: List[asyncio.Task] = []
        # Pin every inner CM so GC can't finalize the parked generator and
        # release the slot for us -- only an explicit __aexit__ should.
        pinned_cms: List = []

        @asynccontextmanager
        async def fake_choose_replica(*args, **kwargs):
            nonlocal reserved_slots
            # __aenter__: reserve a slot.
            with reserved_lock:
                reserved_slots += 1

            # Open the race window deterministically: cancel on the outer loop,
            # then block the router thread until the cancel has propagated -- the
            # caller observes CancelledError before __aexit__ runs.
            cancel_propagated = threading.Event()

            def cancel_and_chain_marker():
                outer_task_holder[0].cancel()
                outer_loop.call_soon(cancel_propagated.set)

            outer_loop.call_soon_threadsafe(cancel_and_chain_marker)
            cancel_propagated.wait()
            try:
                yield fake_selection
            finally:
                # GeneratorExit => released by GC/aclose; else by explicit __aexit__.
                release_path.append(
                    "gc" if isinstance(sys.exc_info()[1], GeneratorExit) else "aexit"
                )
                with reserved_lock:
                    reserved_slots -= 1
                released.set()

        def pinning_factory(*args, **kwargs):
            cm = fake_choose_replica(*args, **kwargs)
            pinned_cms.append(cm)
            return cm

        monkeypatch.setattr(fake_router, "choose_replica", pinning_factory)

        request_metadata = RequestMetadata(
            request_id="test-request-1",
            internal_request_id="test-internal-request-1",
        )

        async def runner():
            async with thread_router.choose_replica(request_metadata):
                pass

        task = asyncio.create_task(runner())
        outer_task_holder.append(task)

        # Cancellation is honored.
        with pytest.raises(asyncio.CancelledError):
            await task

        # Slot released via __aexit__, not GC.
        assert released.wait(timeout=2.0)
        assert release_path == ["aexit"]
        with reserved_lock:
            assert reserved_slots == 0

    @pytest.mark.asyncio
    async def test_finally_shields_cleanup_from_cancellation(
        self,
        setup_router: Tuple[AsyncioRouter, FakeRequestRouter],
        setup_singleton_thread_router: SingletonThreadRouter,
        monkeypatch,
    ):
        """Cancellation during the bridge's finally must not abort cleanup on
        the router loop. The cleanup await is shielded so __aexit__ runs to
        completion even if a cancel arrives while we're awaiting it.
        """
        fake_router, _ = setup_router
        thread_router = setup_singleton_thread_router
        router_loop = thread_router._get_singleton_asyncio_loop(component="unknown")

        fake_selection = ReplicaSelection(
            replica_id="fake-replica",
            node_ip="127.0.0.1",
            port=None,
            node_id="fake-node",
            availability_zone=None,
            replica_metadata={},
            _replica=None,
            _deployment_id=None,
            _request_metadata=None,
            _method_name="",
            _slot_token="",
        )

        aexit_started = threading.Event()
        exit_called = threading.Event()

        async def make_event():
            return asyncio.Event()

        release_event = await asyncio.wrap_future(
            asyncio.run_coroutine_threadsafe(make_event(), router_loop)
        )

        @asynccontextmanager
        async def fake_choose_replica(*args, **kwargs):
            try:
                yield fake_selection
            finally:
                # Signal that the bridge's finally is now awaiting cleanup.
                aexit_started.set()
                # Block until released. If a cancel propagates through to
                # this loop the wait raises and exit_called is never set.
                await release_event.wait()
                exit_called.set()

        monkeypatch.setattr(fake_router, "choose_replica", fake_choose_replica)

        request_metadata = RequestMetadata(
            request_id="test-request-1",
            internal_request_id="test-internal-request-1",
        )

        async def runner():
            async with thread_router.choose_replica(request_metadata):
                pass

        task = asyncio.create_task(runner())

        # Wait until the bridge's finally is awaiting cleanup.
        await asyncio.to_thread(aexit_started.wait)

        # Cancel mid-cleanup. Without shielding, the cancel propagates
        # through wrap_future to the router-loop task and interrupts
        # fake's __aexit__ before it can set exit_called.
        task.cancel()
        await asyncio.sleep(0)

        # Release the gate. With shielding, fake's __aexit__ will now
        # complete and set exit_called.
        router_loop.call_soon_threadsafe(release_event.set)

        with pytest.raises(asyncio.CancelledError):
            await task

        assert exit_called.wait(timeout=2.0), (
            "Cleanup was aborted by cancellation propagating to the router "
            "loop; bridge's finally await must be shielded."
        )

    @pytest.mark.asyncio
    async def test_cancellation_propagation(
        self,
        setup_router: Tuple[AsyncioRouter, FakeRequestRouter],
        setup_singleton_thread_router: SingletonThreadRouter,
    ):
        fake_router, _ = setup_router
        thread_router = setup_singleton_thread_router

        loop = thread_router._get_singleton_asyncio_loop(component="unknown")

        async def init_events():
            return asyncio.Event()

        f = asyncio.run_coroutine_threadsafe(init_events(), loop)
        lock_acquire_event = f.result()
        cancel_block_event = threading.Event()

        fake_replica_result = FakeReplicaResult(
            replica_id=ReplicaID(
                unique_id="test-replica-1", deployment_id=DeploymentID(name="test")
            ),
            is_generator_object=False,
        )

        async def mock_assign_request(*args, **kwargs):
            try:
                await lock_acquire_event.wait()
                return fake_replica_result
            except asyncio.CancelledError:
                cancel_block_event.wait()
                fake_replica_result.cancel()
                raise

        fake_router.assign_request = mock_assign_request

        request_metadata = RequestMetadata(
            request_id="test-request-1",
            internal_request_id="test-internal-request-1",
        )

        assign_request_future = thread_router.assign_request(request_metadata)
        assign_request_future.cancel()

        assert fake_replica_result.cancelled is False

        cancel_block_event.set()

        def release_lock():
            lock_acquire_event.set()

        loop.call_soon_threadsafe(release_lock)

        with pytest.raises(concurrent.futures.CancelledError):
            assign_request_future.result()
        assert assign_request_future.cancelled() is True
        wait_for_condition(lambda: fake_replica_result.cancelled is True)

    @pytest.mark.asyncio
    async def test_replica_result_cancellation(
        self,
        setup_router: Tuple[AsyncioRouter, FakeRequestRouter],
        setup_singleton_thread_router: SingletonThreadRouter,
    ):
        fake_router, _ = setup_router
        thread_router = setup_singleton_thread_router

        fake_replica_result = FakeReplicaResult(
            replica_id=ReplicaID(
                unique_id="test-replica-1", deployment_id=DeploymentID(name="test")
            ),
            is_generator_object=False,
        )

        event = threading.Event()

        async def mock_assign_request(*args, **kwargs):
            # this is a blocking call that will block the asyncio loop
            event.wait()
            return fake_replica_result

        fake_router.assign_request = mock_assign_request

        request_metadata = RequestMetadata(
            request_id="test-request-1",
            internal_request_id="test-internal-request-1",
        )
        assign_request_future = thread_router.assign_request(request_metadata)
        assign_request_future.cancel()

        with pytest.raises(concurrent.futures.CancelledError):
            assign_request_future.result()

        assert assign_request_future.cancelled() is True
        assert fake_replica_result.cancelled is False

        event.set()
        wait_for_condition(lambda: fake_replica_result.cancelled is True)

    @pytest.mark.asyncio
    async def test_assign_request_with_exception(
        self,
        setup_router: Tuple[AsyncioRouter, FakeRequestRouter],
        setup_singleton_thread_router: SingletonThreadRouter,
    ):
        fake_router, _ = setup_router
        thread_router = setup_singleton_thread_router

        async def mock_assign_request(*args, **kwargs):
            raise Exception("test exception")

        fake_router.assign_request = mock_assign_request

        request_metadata = RequestMetadata(
            request_id="test-request-1",
            internal_request_id="test-internal-request-1",
        )
        assign_request_future = thread_router.assign_request(request_metadata)

        assert assign_request_future.exception() is not None
        assert str(assign_request_future.exception()) == "test exception"

    @pytest.mark.asyncio
    async def test_assign_request_with_exception_during_cancellation(
        self,
        setup_router: Tuple[AsyncioRouter, FakeRequestRouter],
        setup_singleton_thread_router: SingletonThreadRouter,
    ):
        fake_router, _ = setup_router
        thread_router = setup_singleton_thread_router

        mock_replica_result = FakeReplicaResult(
            replica_id=ReplicaID(
                unique_id="test-replica-1", deployment_id=DeploymentID(name="test")
            ),
            is_generator_object=False,
        )
        loop = thread_router._get_singleton_asyncio_loop(component="unknown")

        async def init_events():
            return asyncio.Event(), asyncio.Event()

        f = asyncio.run_coroutine_threadsafe(init_events(), loop)
        lock_acquire_event, cancel_block_event = f.result()

        async def mock_assign_request(*args, **kwargs):
            try:
                await lock_acquire_event.wait()
                return mock_replica_result
            except asyncio.CancelledError:
                cancel_block_event.set()
                raise Exception("test exception")

        fake_router.assign_request = mock_assign_request

        request_metadata = RequestMetadata(
            request_id="test-request-1",
            internal_request_id="test-internal-request-1",
        )

        assign_request_future = thread_router.assign_request(request_metadata)
        assign_request_future.cancel()

        async def coro():
            await cancel_block_event.wait()

        asyncio.run_coroutine_threadsafe(coro(), loop).result()

        def release_lock():
            lock_acquire_event.set()

        loop.call_soon_threadsafe(release_lock)

        assert assign_request_future.cancelled() is True
        with pytest.raises(concurrent.futures.CancelledError):
            assign_request_future.exception()


@pytest.mark.asyncio
class TestAsyncioRouterBackoffConfig:
    """Test that backoff config flows from DeploymentConfig to RequestRouter."""

    async def test_update_deployment_config_sets_backoff_params(self):
        """Test that update_deployment_config extracts backoff params from config."""
        fake_request_router = FakeRequestRouter(use_queue_len_cache=False)
        router = AsyncioRouter(
            controller_handle=Mock(),
            deployment_id=DeploymentID(name="test-deployment"),
            handle_id="test-handle-id",
            self_actor_id="test-node-id",
            handle_source=DeploymentHandleSource.UNKNOWN,
            event_loop=get_or_create_event_loop(),
            enable_strict_max_ongoing_requests=False,
            request_router=fake_request_router,
            node_id="test-node-id",
            availability_zone="test-az",
            prefer_local_node_routing=False,
            _request_router_initialized_event=asyncio.Event(),
        )

        # Create a DeploymentConfig with custom backoff params
        custom_initial_backoff = 0.15
        custom_multiplier = 4
        custom_max_backoff = 2.5

        deployment_config = DeploymentConfig.from_default(
            request_router_config=RequestRouterConfig(
                initial_backoff_s=custom_initial_backoff,
                backoff_multiplier=custom_multiplier,
                max_backoff_s=custom_max_backoff,
            )
        )

        # Update the router with the config
        router.update_deployment_config(deployment_config)

        # Verify the backoff params were stored on the router
        assert router._initial_backoff_s == custom_initial_backoff
        assert router._backoff_multiplier == custom_multiplier
        assert router._max_backoff_s == custom_max_backoff

        # Verify the backoff params were propagated to the request router
        assert fake_request_router.initial_backoff_s == custom_initial_backoff
        assert fake_request_router.backoff_multiplier == custom_multiplier
        assert fake_request_router.max_backoff_s == custom_max_backoff


class TestOnRequestCompleted:
    """Tests for the on_request_completed hook introduced in this branch.

    Verifies that on_request_completed is called on the request router:
    - Via the done-callback when a request completes normally.
    - Via the finally block when a request fails before the callback is registered.
    """

    async def test_on_request_completed_called_via_done_callback(
        self,
        setup_router: Tuple[AsyncioRouter, FakeRequestRouter],
    ):
        """After a normal (no-rejection) request, firing the ReplicaResult's
        done callback should invoke on_request_completed on the router."""
        router, fake_request_router = setup_router

        r1_id = ReplicaID(
            unique_id="test-replica-1", deployment_id=DeploymentID(name="test")
        )
        fake_request_router.set_replica_to_return(FakeReplica(r1_id))

        metadata = RequestMetadata(
            request_id="req-1",
            internal_request_id="internal-req-1",
        )
        result = await router.assign_request(metadata)

        assert len(fake_request_router.completed_requests) == 0

        # Simulate the request finishing.
        result.fire_done_callbacks(result=None)
        await asyncio.sleep(0)

        assert len(fake_request_router.completed_requests) == 1
        completed_replica_id, completed_req_id = fake_request_router.completed_requests[
            0
        ]
        assert completed_replica_id == r1_id
        assert completed_req_id == "internal-req-1"

    @pytest.mark.parametrize(
        "setup_router",
        [{"enable_strict_max_ongoing_requests": True}],
        indirect=True,
    )
    async def test_on_request_completed_called_via_finally_on_actor_died(
        self,
        setup_router: Tuple[AsyncioRouter, FakeRequestRouter],
    ):
        """When try_send_request raises ActorDiedError (before the callback is
        registered), the finally block should call on_request_completed."""
        router, fake_request_router = setup_router
        d_id = DeploymentID(name="test")
        r1_id = ReplicaID(unique_id="r1", deployment_id=d_id)
        r2_id = ReplicaID(unique_id="r2", deployment_id=d_id)

        fake_request_router.set_replica_to_return(
            FakeReplica(r1_id, error=ActorDiedError())
        )
        fake_request_router.set_replica_to_return_on_retry(
            FakeReplica(
                r2_id,
                queue_len_info=ReplicaQueueLengthInfo(
                    accepted=True, num_ongoing_requests=5
                ),
            )
        )

        await router.assign_request(dummy_request_metadata())

        # The finally block should have called on_request_completed for r1.
        r1_completions = [
            (rid, req_id)
            for rid, req_id in fake_request_router.completed_requests
            if rid == r1_id
        ]
        assert len(r1_completions) == 1

    @pytest.mark.parametrize(
        "setup_router",
        [{"enable_strict_max_ongoing_requests": True}],
        indirect=True,
    )
    async def test_on_request_completed_called_via_finally_on_actor_unavailable(
        self,
        setup_router: Tuple[AsyncioRouter, FakeRequestRouter],
    ):
        """When try_send_request raises ActorUnavailableError, the finally
        block should call on_request_completed."""
        router, fake_request_router = setup_router
        d_id = DeploymentID(name="test")
        r1_id = ReplicaID(unique_id="r1", deployment_id=d_id)
        r2_id = ReplicaID(unique_id="r2", deployment_id=d_id)

        fake_request_router.set_replica_to_return(
            FakeReplica(
                r1_id,
                error=ActorUnavailableError(error_message="unavailable", actor_id=None),
            )
        )
        fake_request_router.set_replica_to_return_on_retry(
            FakeReplica(
                r2_id,
                queue_len_info=ReplicaQueueLengthInfo(
                    accepted=True, num_ongoing_requests=5
                ),
            )
        )

        await router.assign_request(dummy_request_metadata())

        r1_completions = [
            (rid, req_id)
            for rid, req_id in fake_request_router.completed_requests
            if rid == r1_id
        ]
        assert len(r1_completions) == 1

    async def test_multiple_requests_each_trigger_on_request_completed(
        self,
        setup_router: Tuple[AsyncioRouter, FakeRequestRouter],
    ):
        """Multiple in-flight requests should each trigger their own
        on_request_completed call when they finish."""
        router, fake_request_router = setup_router

        r1_id = ReplicaID(
            unique_id="test-replica-1", deployment_id=DeploymentID(name="test")
        )
        fake_request_router.set_replica_to_return(FakeReplica(r1_id))

        results = []
        for i in range(5):
            metadata = RequestMetadata(
                request_id=f"req-{i}",
                internal_request_id=f"internal-req-{i}",
            )
            result = await router.assign_request(metadata)
            results.append(result)

        assert len(fake_request_router.completed_requests) == 0

        # Complete them one at a time.
        for i, result in enumerate(results):
            result.fire_done_callbacks(result=None)
            await asyncio.sleep(0)
            assert len(fake_request_router.completed_requests) == i + 1

        completed_ids = {req_id for _, req_id in fake_request_router.completed_requests}
        expected_ids = {f"internal-req-{i}" for i in range(5)}
        assert completed_ids == expected_ids


@pytest.mark.asyncio
class TestCustomRequestRouterAPIs:
    """Tests for the new RequestRouter APIs introduced in this branch:
    - supports_rejection_protocol property
    - _compute_backoff_s / _backoff helpers
    """

    def test_supports_rejection_protocol_defaults_to_true(self):
        """The base RequestRouter.supports_rejection_protocol should be True."""

        class MinimalRouter(RequestRouter):
            async def choose_replicas(self, candidate_replicas, pending_request=None):
                return [candidate_replicas]

        r = MinimalRouter(
            deployment_id=DeploymentID(name="test"),
            handle_source=DeploymentHandleSource.UNKNOWN,
            create_replica_wrapper_func=lambda ri: FakeReplica(ri),
        )
        assert r.supports_rejection_protocol is True

    def test_supports_rejection_protocol_can_be_overridden(self):
        """A subclass can override supports_rejection_protocol to False."""

        class NoRejectionRouter(RequestRouter):
            async def choose_replicas(self, candidate_replicas, pending_request=None):
                return [candidate_replicas]

            @property
            def supports_rejection_protocol(self) -> bool:
                return False

        r = NoRejectionRouter(
            deployment_id=DeploymentID(name="test"),
            handle_source=DeploymentHandleSource.UNKNOWN,
            create_replica_wrapper_func=lambda ri: FakeReplica(ri),
        )
        assert r.supports_rejection_protocol is False

    def test_compute_backoff_s_uses_class_parameters(self):
        """_compute_backoff_s should respect the class-level backoff params."""

        class MinimalRouter(RequestRouter):
            async def choose_replicas(self, candidate_replicas, pending_request=None):
                return [candidate_replicas]

        r = MinimalRouter(
            deployment_id=DeploymentID(name="test"),
            handle_source=DeploymentHandleSource.UNKNOWN,
            create_replica_wrapper_func=lambda ri: FakeReplica(ri),
        )

        b0 = r._compute_backoff_s(0)
        b1 = r._compute_backoff_s(1)
        assert b0 == r.initial_backoff_s
        assert b1 == pytest.approx(r.initial_backoff_s * r.backoff_multiplier)

    def test_compute_backoff_s_with_custom_multiplier(self):
        """A subclass can override backoff_multiplier and _compute_backoff_s
        will respect it."""

        class CustomBackoffRouter(RequestRouter):
            async def choose_replicas(self, candidate_replicas, pending_request=None):
                return [candidate_replicas]

        r = CustomBackoffRouter(
            deployment_id=DeploymentID(name="test"),
            handle_source=DeploymentHandleSource.UNKNOWN,
            create_replica_wrapper_func=lambda ri: FakeReplica(ri),
            backoff_multiplier=1.5,
        )

        b0 = r._compute_backoff_s(0)
        b1 = r._compute_backoff_s(1)
        b2 = r._compute_backoff_s(2)

        assert b0 == r.initial_backoff_s
        assert b1 == pytest.approx(r.initial_backoff_s * 1.5)
        assert b2 == pytest.approx(r.initial_backoff_s * 1.5**2)

    def test_compute_backoff_s_capped_at_max(self):
        """Backoff should never exceed max_backoff_s."""

        class MinimalRouter(RequestRouter):
            async def choose_replicas(self, candidate_replicas, pending_request=None):
                return [candidate_replicas]

        r = MinimalRouter(
            deployment_id=DeploymentID(name="test"),
            handle_source=DeploymentHandleSource.UNKNOWN,
            create_replica_wrapper_func=lambda ri: FakeReplica(ri),
        )

        for attempt in [50, 100, 1000]:
            assert r._compute_backoff_s(attempt) <= r.max_backoff_s

    async def test_backoff_is_awaitable(self):
        """_backoff should be an awaitable coroutine."""

        class MinimalRouter(RequestRouter):
            async def choose_replicas(self, candidate_replicas, pending_request=None):
                return [candidate_replicas]

        r = MinimalRouter(
            deployment_id=DeploymentID(name="test"),
            handle_source=DeploymentHandleSource.UNKNOWN,
            create_replica_wrapper_func=lambda ri: FakeReplica(ri),
        )
        # Should complete without error.
        await r._backoff(0)


if __name__ == "__main__":
    sys.exit(pytest.main(["-v", "-s", __file__]))