[feat] Improve API: streaming router (multi-replica load balancer + ws proxy) (#1286)

Co-authored-by: Junda (David) Su <90978028+Davids048@users.noreply.github.com>
Co-authored-by: Matthew Noto <99706358+RandNMR73@users.noreply.github.com>
Co-authored-by: XOR-op <17672363+XOR-op@users.noreply.github.com>
Co-authored-by: Zhang Peiyuan <42993249+jzhang38@users.noreply.github.com>

fastvideo/entrypoints/cli/main.py

@@ -3,6 +3,8 @@
 from fastvideo.entrypoints.cli.cli_types import CLISubcommand
 from fastvideo.entrypoints.cli.generate import cmd_init as generate_cmd_init
 from fastvideo.utils import FlexibleArgumentParser
+from fastvideo.entrypoints.cli.router_serve import (
+    cmd_init as router_serve_cmd_init, )
 from fastvideo.entrypoints.cli.serve import cmd_init as serve_cmd_init
 from fastvideo.entrypoints.cli.bench import cmd_init as bench_cmd_init
 
@@ -12,6 +14,7 @@ def cmd_init() -> list[CLISubcommand]:
     commands = []
     commands.extend(generate_cmd_init())
     commands.extend(serve_cmd_init())
+    commands.extend(router_serve_cmd_init())
     commands.extend(bench_cmd_init())
     return commands
 

fastvideo/entrypoints/cli/router_serve.py

@@ -0,0 +1,115 @@
+# SPDX-License-Identifier: Apache-2.0
+"""``fastvideo router-serve`` CLI subcommand.
+
+Launches the streaming router from a YAML config. Separate from
+``fastvideo serve`` because the router is an orthogonal process: it
+fronts one or more running servers rather than hosting a generator
+itself.
+"""
+from __future__ import annotations
+
+import argparse
+import os
+from typing import cast
+
+from fastvideo.api.parser import load_raw_config
+from fastvideo.entrypoints.cli.cli_types import CLISubcommand
+from fastvideo.entrypoints.streaming.router.config import (
+    ReplicaEndpoint,
+    RouterConfig,
+)
+from fastvideo.logger import init_logger
+from fastvideo.utils import FlexibleArgumentParser
+
+logger = init_logger(__name__)
+
+
+class RouterServeSubcommand(CLISubcommand):
+    """Start the multi-replica WebSocket router."""
+
+    def __init__(self) -> None:
+        self.name = "router-serve"
+        super().__init__()
+
+    def cmd(self, args: argparse.Namespace) -> None:
+        config = _load_router_config(args.config)
+        logger.info(
+            "router listening on %s:%d (%d replicas, %d primary)",
+            config.host,
+            config.port,
+            len(config.replicas),
+            sum(1 for r in config.replicas if r.primary),
+        )
+        from fastvideo.entrypoints.streaming.router.main import run_router
+
+        run_router(config)
+
+    def validate(self, args: argparse.Namespace) -> None:
+        if not args.config:
+            raise ValueError("fastvideo router-serve requires --config PATH")
+        if not os.path.exists(args.config):
+            raise ValueError(f"Router config file not found: {args.config}")
+
+    def subparser_init(
+        self,
+        subparsers: argparse._SubParsersAction,
+    ) -> FlexibleArgumentParser:
+        parser = subparsers.add_parser(
+            "router-serve",
+            help="Start the streaming router (multi-replica load balancer)",
+            usage="fastvideo router-serve --config ROUTER_CONFIG",
+        )
+        parser.add_argument(
+            "--config",
+            type=str,
+            default="",
+            required=False,
+            help="Path to a YAML/JSON router config. Required.",
+        )
+        return cast(FlexibleArgumentParser, parser)
+
+
+def _load_router_config(path: str) -> RouterConfig:
+    raw = load_raw_config(path)
+    router_raw = raw.get("router") if isinstance(raw, dict) else None
+    if not isinstance(router_raw, dict):
+        raise ValueError(f"Router config {path!r} must have a top-level `router:` block")
+
+    replicas_raw = router_raw.get("replicas", [])
+    if not isinstance(replicas_raw, list):
+        raise ValueError(f"router.replicas must be a list, got {type(replicas_raw).__name__}")
+    replicas = []
+    for i, r in enumerate(replicas_raw):
+        if not isinstance(r, dict):
+            raise ValueError(f"router.replicas[{i}] must be a mapping, got {type(r).__name__}")
+        url = r.get("url")
+        if not url:
+            raise ValueError(f"router.replicas[{i}] is missing required key 'url'")
+        replicas.append(
+            ReplicaEndpoint(
+                url=url,
+                name=r.get("name"),
+                primary=bool(r.get("primary", False)),
+                weight=float(r.get("weight", 1.0)),
+            ))
+    if not replicas:
+        raise ValueError("Router config must list at least one replica under `router.replicas`")
+
+    health_check = router_raw.get("health_check") or {}
+    return RouterConfig(
+        host=str(router_raw.get("host", "0.0.0.0")),
+        port=int(router_raw.get("port", 9000)),
+        replicas=replicas,
+        health_check_path=str(health_check.get("path", "/health")),
+        health_check_interval_seconds=float(health_check.get("interval_seconds", 5.0)),
+        health_check_timeout_seconds=float(health_check.get("timeout_seconds", 2.0)),
+        failure_threshold=int(health_check.get("failure_threshold", 3)),
+        recovery_threshold=int(health_check.get("recovery_threshold", 2)),
+    )
+
+
+def cmd_init() -> list[CLISubcommand]:
+    return [RouterServeSubcommand()]
+
+
+__all__ = ["RouterServeSubcommand", "cmd_init"]

fastvideo/entrypoints/streaming/router/__init__.py

@@ -0,0 +1,27 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Multi-replica load balancer + WebSocket proxy for the streaming server.
+
+Sits in front of one-or-more streaming-server replicas and forwards
+WebSocket sessions to a healthy primary, with failover to secondaries.
+Kept in-repo under ``fastvideo/entrypoints/streaming/router/`` per the
+PR plan's default; the alternative (separate package) is an open
+question deferred to review.
+"""
+from fastvideo.entrypoints.streaming.router.registry import (
+    Replica,
+    ReplicaHealth,
+    ReplicaRegistry,
+    ReplicaStatus,
+)
+from fastvideo.entrypoints.streaming.router.config import RouterConfig
+from fastvideo.entrypoints.streaming.router.main import build_router_app, run_router
+
+__all__ = [
+    "Replica",
+    "ReplicaHealth",
+    "ReplicaRegistry",
+    "ReplicaStatus",
+    "RouterConfig",
+    "build_router_app",
+    "run_router",
+]

fastvideo/entrypoints/streaming/router/config.py

@@ -0,0 +1,88 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Typed router configuration."""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from urllib.parse import urlparse
+
+
+@dataclass
+class ReplicaEndpoint:
+    """One backend replica the router can route to."""
+
+    url: str
+    """HTTP base URL, e.g. ``http://host:8000``. WebSocket URL is
+    derived automatically by replacing the scheme."""
+    name: str | None = None
+    primary: bool = False
+    """``True`` = prefer this replica over others in steady state."""
+    weight: float = 1.0
+
+
+@dataclass
+class RouterConfig:
+    """Typed router config loaded from a YAML file.
+
+    Example::
+
+        router:
+          host: 0.0.0.0
+          port: 9000
+          replicas:
+            - url: http://streamer-a:8000
+              primary: true
+            - url: http://streamer-b:8000
+          health_check:
+            path: /health
+            interval_seconds: 5
+            failure_threshold: 3
+
+    Validation runs in ``__post_init__``: empty replicas, non-positive
+    intervals/timeouts, thresholds < 1, non-http(s) URLs, and more than
+    one primary all raise ``ValueError`` so misconfigurations surface at
+    load time rather than as confusing runtime failures.
+    """
+
+    host: str = "0.0.0.0"
+    port: int = 9000
+    replicas: list[ReplicaEndpoint] = field(default_factory=list)
+    health_check_path: str = "/health"
+    health_check_interval_seconds: float = 5.0
+    health_check_timeout_seconds: float = 2.0
+    failure_threshold: int = 3
+    recovery_threshold: int = 2
+
+    def __post_init__(self) -> None:
+        if not self.replicas:
+            raise ValueError("RouterConfig.replicas must list at least one replica")
+        if self.health_check_interval_seconds <= 0:
+            raise ValueError(f"health_check_interval_seconds must be > 0, got {self.health_check_interval_seconds}")
+        if self.health_check_timeout_seconds <= 0:
+            raise ValueError(f"health_check_timeout_seconds must be > 0, got {self.health_check_timeout_seconds}")
+        if self.failure_threshold < 1:
+            raise ValueError(f"failure_threshold must be >= 1, got {self.failure_threshold}")
+        if self.recovery_threshold < 1:
+            raise ValueError(f"recovery_threshold must be >= 1, got {self.recovery_threshold}")
+        seen_urls: set[str] = set()
+        for replica in self.replicas:
+            if not replica.url.startswith(("http://", "https://")):
+                raise ValueError(f"ReplicaEndpoint.url must start with http:// or https://, got {replica.url!r}")
+            parsed = urlparse(replica.url)
+            if parsed.path not in ("", "/"):
+                raise ValueError(f"ReplicaEndpoint.url must be a base host[:port] URL without a path; "
+                                 f"got {replica.url!r} with path {parsed.path!r}. The router appends "
+                                 "`/health` and `/v1/stream` itself.")
+            if parsed.query or parsed.fragment:
+                raise ValueError(f"ReplicaEndpoint.url must not include query/fragment; got {replica.url!r}")
+            if replica.url in seen_urls:
+                raise ValueError(f"Duplicate ReplicaEndpoint.url {replica.url!r}; "
+                                 "router selection keys by URL so duplicates would silently collapse")
+            seen_urls.add(replica.url)
+        primaries = sum(1 for r in self.replicas if r.primary)
+        if primaries > 1:
+            raise ValueError(f"RouterConfig allows at most one primary replica; got {primaries}. "
+                             "Multi-primary load distribution is deferred — promote one replica to "
+                             "primary and treat the rest as secondaries.")
+
+
+__all__ = ["ReplicaEndpoint", "RouterConfig"]

fastvideo/entrypoints/streaming/router/main.py

@@ -0,0 +1,218 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Router FastAPI entry point.
+
+Exposes the same ``/v1/stream`` WebSocket path the backend servers do,
+accepts a client, picks a healthy replica from the registry, and
+proxies frames bidirectionally.
+
+PR 7.9 ships the minimum-viable shape: explicit replica list, single
+primary, JSON + binary passthrough in both directions, and a
+``/status`` endpoint for operators. Sticky-session routing (so a
+reconnect lands on the same backend) is left for a follow-up.
+"""
+from __future__ import annotations
+
+import asyncio
+import contextlib
+from dataclasses import dataclass
+
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect
+from fastapi.responses import JSONResponse
+
+from fastvideo.entrypoints.streaming.router.config import RouterConfig
+from fastvideo.entrypoints.streaming.router.registry import (
+    ReplicaRegistry,
+    run_health_check_loop,
+)
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class _RouterState:
+    config: RouterConfig
+    registry: ReplicaRegistry
+    stop_event: asyncio.Event
+    health_task: asyncio.Task | None = None
+
+
+def build_router_app(
+    config: RouterConfig,
+    *,
+    registry: ReplicaRegistry | None = None,
+) -> FastAPI:
+    """Build the router FastAPI app.
+
+    ``registry`` can be injected for tests; defaults to one built from
+    ``config.replicas``.
+    """
+    registry = registry or ReplicaRegistry(config.replicas)
+    state = _RouterState(
+        config=config,
+        registry=registry,
+        stop_event=asyncio.Event(),
+    )
+
+    @contextlib.asynccontextmanager
+    async def _lifespan(_app: FastAPI):
+        state.health_task = asyncio.create_task(
+            run_health_check_loop(
+                registry=state.registry,
+                config=state.config,
+                stop_event=state.stop_event,
+            ))
+        try:
+            yield
+        finally:
+            state.stop_event.set()
+            if state.health_task is not None:
+                with contextlib.suppress(asyncio.CancelledError):
+                    await state.health_task
+
+    app = FastAPI(title="FastVideo Streaming Router", lifespan=_lifespan)
+
+    @app.get("/status")
+    async def _status() -> JSONResponse:
+        return JSONResponse({
+            "replicas": [{
+                "url": r.url,
+                "primary": r.primary,
+                "status": r.health.status.value,
+                "last_ok_at": r.health.last_ok_at,
+                "last_latency_ms": r.health.last_latency_ms,
+                "consecutive_failures": r.health.consecutive_failures,
+            } for r in state.registry.all()],
+        })
+
+    @app.websocket("/v1/stream")
+    async def _proxy(websocket: WebSocket) -> None:
+        await websocket.accept()
+        replica = state.registry.select()
+        if replica is None:
+            await websocket.send_json({
+                "type": "error",
+                "code": "gpu_unavailable",
+                "message": "router: no healthy replica available",
+                "retryable": True,
+            })
+            await websocket.close(code=1013, reason="no_healthy_replica")
+            return
+
+        ws_url = _websocket_url_for(replica.url)
+        try:
+            await _bridge_session(websocket, ws_url)
+        except WebSocketDisconnect:
+            logger.info("router: client disconnected")
+        except Exception as exc:
+            logger.exception("router: bridge failed: %s", exc)
+            with contextlib.suppress(RuntimeError):
+                await websocket.send_json({
+                    "type": "error",
+                    "code": "worker_failed",
+                    "message": f"router bridge failed: {exc}",
+                    "retryable": True,
+                })
+            with contextlib.suppress(RuntimeError):
+                await websocket.close(code=1011)
+
+    app.state.router_state = state
+    return app
+
+
+def run_router(config: RouterConfig) -> None:  # pragma: no cover - CLI
+    import uvicorn
+
+    app = build_router_app(config)
+    uvicorn.run(app, host=config.host, port=config.port)
+
+
+async def _bridge_session(
+    client_ws: WebSocket,
+    backend_ws_url: str,
+) -> None:
+    """Connect to backend and shuttle messages in both directions.
+
+    Uses ``websockets`` for the backend side; imported lazily to keep
+    the router's import graph small for users who only want the server.
+
+    Cancellation: when either direction completes (client disconnect,
+    backend close, exception), the other is cancelled explicitly and
+    both are drained before returning. Unexpected exceptions from the
+    direction that completed first are re-raised; normal disconnect
+    paths (``WebSocketDisconnect``, ``ConnectionClosed``,
+    ``CancelledError``) are swallowed.
+    """
+    try:
+        import websockets
+    except ImportError as exc:  # pragma: no cover - optional extra
+        raise RuntimeError("router requires the `websockets` package for backend proxying") from exc
+
+    async with websockets.connect(backend_ws_url + "/v1/stream") as backend_ws:
+        c2b = asyncio.create_task(_forward_client_to_backend(client_ws, backend_ws))
+        b2c = asyncio.create_task(_forward_backend_to_client(backend_ws, client_ws))
+        try:
+            done, _pending = await asyncio.wait(
+                {c2b, b2c},
+                return_when=asyncio.FIRST_COMPLETED,
+            )
+        finally:
+            for task in (c2b, b2c):
+                if not task.done():
+                    task.cancel()
+            await asyncio.gather(c2b, b2c, return_exceptions=True)
+        for task in done:
+            task_exc = task.exception()
+            if task_exc is not None and not _is_normal_disconnect(task_exc):
+                raise task_exc
+
+
+def _is_normal_disconnect(exc: BaseException) -> bool:
+    """Whether ``exc`` is a routine WebSocket teardown vs a real bridge fault."""
+    if isinstance(exc, asyncio.CancelledError | WebSocketDisconnect):
+        return True
+    name = type(exc).__name__
+    # websockets.exceptions.ConnectionClosed{,OK,Error} all subclass
+    # WebSocketException; check by name to avoid the lazy-import dance.
+    return name.startswith("ConnectionClosed")
+
+
+async def _forward_client_to_backend(client_ws: WebSocket, backend_ws) -> None:
+    try:
+        while True:
+            msg = await client_ws.receive()
+            if msg.get("type") == "websocket.disconnect":
+                break
+            if "text" in msg and msg["text"] is not None:
+                await backend_ws.send(msg["text"])
+            elif "bytes" in msg and msg["bytes"] is not None:
+                await backend_ws.send(msg["bytes"])
+    finally:
+        with contextlib.suppress(Exception):
+            await backend_ws.close()
+
+
+async def _forward_backend_to_client(backend_ws, client_ws: WebSocket) -> None:
+    try:
+        async for frame in backend_ws:
+            if isinstance(frame, bytes):
+                await client_ws.send_bytes(frame)
+            else:
+                await client_ws.send_text(frame)
+    finally:
+        with contextlib.suppress(Exception):
+            await client_ws.close()
+
+
+def _websocket_url_for(http_url: str) -> str:
+    if http_url.startswith("https://"):
+        return "wss://" + http_url[len("https://"):]
+    if http_url.startswith("http://"):
+        return "ws://" + http_url[len("http://"):]
+    return http_url
+
+
+__all__ = [
+    "build_router_app",
+    "run_router",
+]

fastvideo/entrypoints/streaming/router/registry.py

@@ -0,0 +1,268 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Replica registry + health-check loop.
+
+The registry tracks the set of known backend replicas and their live
+health. The router consults it for "pick a backend for this session"
+decisions and a background task updates it from periodic HTTP probes.
+
+State machine per replica::
+
+    HEALTHY ──(N consecutive failures)──▶ UNHEALTHY
+       ▲                                     │
+       └──────(M consecutive successes)──────┘
+
+Where N = :attr:`RouterConfig.failure_threshold` and
+M = :attr:`RouterConfig.recovery_threshold`.
+"""
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import enum
+import time
+from collections.abc import AsyncIterator, Awaitable, Callable
+from dataclasses import dataclass, field
+from typing import Any
+
+from fastvideo.entrypoints.streaming.router.config import (
+    ReplicaEndpoint,
+    RouterConfig,
+)
+from fastvideo.logger import init_logger
+
+HttpProbe = Any
+"""Structural alias for health-probe callables. Concrete signature is
+``async def __call__(url: str, *, timeout: float) -> tuple[float,
+str | None]``; typing.Callable cannot express keyword-only parameters,
+so duck-typing is the pragmatic compromise."""
+
+logger = init_logger(__name__)
+
+
+class ReplicaStatus(enum.Enum):
+    UNKNOWN = "unknown"
+    HEALTHY = "healthy"
+    UNHEALTHY = "unhealthy"
+
+
+@dataclass
+class ReplicaHealth:
+    status: ReplicaStatus = ReplicaStatus.UNKNOWN
+    last_ok_at: float | None = None
+    last_failure_at: float | None = None
+    consecutive_failures: int = 0
+    consecutive_successes: int = 0
+    last_latency_ms: float | None = None
+
+
+@dataclass
+class Replica:
+    endpoint: ReplicaEndpoint
+    health: ReplicaHealth = field(default_factory=ReplicaHealth)
+
+    @property
+    def url(self) -> str:
+        return self.endpoint.url
+
+    @property
+    def primary(self) -> bool:
+        return self.endpoint.primary
+
+    @property
+    def is_healthy(self) -> bool:
+        return self.health.status is ReplicaStatus.HEALTHY
+
+
+class ReplicaRegistry:
+    """Stateful map of replica URL → :class:`Replica`.
+
+    Selection favors primary replicas when healthy; otherwise the first
+    healthy non-primary is returned. When none are healthy, the
+    registry returns ``None`` so the router can reject incoming
+    sessions with ``gpu_unavailable``.
+    """
+
+    def __init__(self, replicas: list[ReplicaEndpoint]) -> None:
+        if not replicas:
+            raise ValueError("ReplicaRegistry requires at least one replica")
+        self._replicas: dict[str, Replica] = {endpoint.url: Replica(endpoint=endpoint) for endpoint in replicas}
+        self._lock = asyncio.Lock()
+
+    def all(self) -> list[Replica]:
+        return list(self._replicas.values())
+
+    def get(self, url: str) -> Replica | None:
+        return self._replicas.get(url)
+
+    def primaries(self) -> list[Replica]:
+        return [r for r in self._replicas.values() if r.primary]
+
+    def select(self) -> Replica | None:
+        """Pick the best healthy replica.
+
+        Priority order:
+
+        1. The first healthy primary (insertion order).
+        2. The first healthy non-primary (insertion order).
+        3. ``None`` when nothing is healthy.
+
+        This MVP picks the first match within each tier; it does NOT
+        load-balance across multiple healthy replicas of the same tier.
+        Round-robin and weighted distribution are deferred until a real
+        N-way active deployment exists.
+        """
+        healthy_primaries = [r for r in self._replicas.values() if r.primary and r.is_healthy]
+        if healthy_primaries:
+            return healthy_primaries[0]
+        healthy = [r for r in self._replicas.values() if r.is_healthy]
+        if healthy:
+            return healthy[0]
+        return None
+
+    async def record_success(
+        self,
+        replica: Replica,
+        *,
+        recovery_threshold: int,
+        latency_ms: float,
+    ) -> None:
+        async with self._lock:
+            h = replica.health
+            h.last_ok_at = time.time()
+            h.last_latency_ms = latency_ms
+            h.consecutive_failures = 0
+            h.consecutive_successes += 1
+            # State machine: UNKNOWN -> HEALTHY is immediate; only the
+            # UNHEALTHY -> HEALTHY transition is gated by recovery_threshold.
+            if h.status is ReplicaStatus.UNKNOWN:
+                logger.info("router: replica %s initial probe ok, marking HEALTHY", replica.url)
+                h.status = ReplicaStatus.HEALTHY
+                h.consecutive_successes = 0
+            elif (h.status is ReplicaStatus.UNHEALTHY and h.consecutive_successes >= recovery_threshold):
+                logger.info("router: replica %s recovered to HEALTHY after %d successes", replica.url,
+                            h.consecutive_successes)
+                h.status = ReplicaStatus.HEALTHY
+                h.consecutive_successes = 0
+
+    async def record_failure(
+        self,
+        replica: Replica,
+        *,
+        failure_threshold: int,
+        reason: str,
+    ) -> None:
+        async with self._lock:
+            h = replica.health
+            h.last_failure_at = time.time()
+            h.consecutive_successes = 0
+            h.consecutive_failures += 1
+            if (h.status is not ReplicaStatus.UNHEALTHY and h.consecutive_failures >= failure_threshold):
+                logger.warning("router: replica %s marked UNHEALTHY after %d failures: %s", replica.url,
+                               h.consecutive_failures, reason)
+                h.status = ReplicaStatus.UNHEALTHY
+
+
+async def run_health_check_loop(
+    registry: ReplicaRegistry,
+    config: RouterConfig,
+    *,
+    stop_event: asyncio.Event,
+    http_get: HttpProbe | None = None,
+) -> None:
+    """Poll all replicas' health endpoints in parallel on a fixed interval.
+
+    ``http_get`` is pluggable so unit tests can inject a deterministic
+    probe without hitting the network. The default builds a single
+    ``httpx.AsyncClient`` shared across the loop's lifetime so the
+    common case (steady polling against a stable replica set) reuses
+    TCP/TLS connections instead of paying handshake cost per probe.
+
+    Probes within one polling cycle run concurrently via ``asyncio.gather``
+    so a slow replica doesn't push the cycle past
+    ``health_check_interval_seconds``.
+    """
+    if http_get is not None:
+        await _run_loop(registry, config, stop_event, http_get)
+        return
+    async with _build_default_probe(config) as probe:
+        await _run_loop(registry, config, stop_event, probe)
+
+
+async def _run_loop(
+    registry: ReplicaRegistry,
+    config: RouterConfig,
+    stop_event: asyncio.Event,
+    http_get: Callable[..., Awaitable[tuple[float, str | None]]],
+) -> None:
+    while not stop_event.is_set():
+        replicas = registry.all()
+        results = await asyncio.gather(
+            *[
+                http_get(replica.url + config.health_check_path, timeout=config.health_check_timeout_seconds)
+                for replica in replicas
+            ],
+            return_exceptions=True,
+        )
+        for replica, result in zip(replicas, results, strict=True):
+            if isinstance(result, BaseException):
+                await registry.record_failure(
+                    replica,
+                    failure_threshold=config.failure_threshold,
+                    reason=f"{type(result).__name__}: {result}",
+                )
+                continue
+            status_ms, error = result
+            if error is None:
+                await registry.record_success(
+                    replica,
+                    recovery_threshold=config.recovery_threshold,
+                    latency_ms=status_ms,
+                )
+            else:
+                await registry.record_failure(
+                    replica,
+                    failure_threshold=config.failure_threshold,
+                    reason=error,
+                )
+        try:
+            await asyncio.wait_for(
+                stop_event.wait(),
+                timeout=config.health_check_interval_seconds,
+            )
+        except asyncio.TimeoutError:
+            continue
+
+
+@contextlib.asynccontextmanager
+async def _build_default_probe(
+    config: RouterConfig, ) -> AsyncIterator[Callable[..., Awaitable[tuple[float, str | None]]]]:
+    try:
+        import httpx
+    except ImportError as exc:  # pragma: no cover - optional extra
+        raise RuntimeError("router health checks require httpx; install with "
+                           "`pip install fastvideo[streaming]` or `pip install httpx`") from exc
+
+    async with httpx.AsyncClient(timeout=config.health_check_timeout_seconds) as client:
+
+        async def probe(url: str, *, timeout: float) -> tuple[float, str | None]:
+            start = time.perf_counter()
+            try:
+                response = await client.get(url, timeout=timeout)
+            except Exception as exc:
+                return 0.0, f"{type(exc).__name__}: {exc}"
+            latency_ms = (time.perf_counter() - start) * 1000.0
+            if response.status_code >= 400:
+                return latency_ms, f"HTTP {response.status_code}"
+            return latency_ms, None
+
+        yield probe
+
+
+__all__ = [
+    "HttpProbe",
+    "Replica",
+    "ReplicaHealth",
+    "ReplicaRegistry",
+    "ReplicaStatus",
+    "run_health_check_loop",
+]

fastvideo/tests/entrypoints/streaming/test_router.py

@@ -0,0 +1,279 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Router tests — registry semantics + health loop behavior.
+
+Avoids real WebSocket proxying (the bridge requires the ``websockets``
+package and two running uvicorn processes); test_server.py already
+covers the end-to-end WS protocol against a direct backend.
+"""
+from __future__ import annotations
+
+import asyncio
+
+import pytest
+
+from fastvideo.entrypoints.streaming.router.config import (
+    ReplicaEndpoint,
+    RouterConfig,
+)
+from fastvideo.entrypoints.streaming.router.registry import (
+    ReplicaRegistry,
+    ReplicaStatus,
+    run_health_check_loop,
+)
+
+
+def _registry(
+    *,
+    num_primary: int = 1,
+    num_secondary: int = 1,
+) -> ReplicaRegistry:
+    replicas = [
+        ReplicaEndpoint(
+            url=f"http://primary-{i}:8000",
+            primary=True,
+        )
+        for i in range(num_primary)
+    ] + [
+        ReplicaEndpoint(
+            url=f"http://secondary-{i}:8000",
+            primary=False,
+        )
+        for i in range(num_secondary)
+    ]
+    return ReplicaRegistry(replicas)
+
+
+class TestReplicaRegistry:
+
+    def test_requires_replicas(self):
+        with pytest.raises(ValueError):
+            ReplicaRegistry([])
+
+    def test_select_none_when_all_unknown(self):
+        assert _registry().select() is None
+
+    def test_select_prefers_healthy_primary(self):
+        reg = _registry()
+
+        async def promote():
+            primary = reg.primaries()[0]
+            await reg.record_success(
+                primary, recovery_threshold=1, latency_ms=1.0)
+            # Also mark secondary healthy — primary should still win.
+            sec = next(r for r in reg.all() if not r.primary)
+            await reg.record_success(
+                sec, recovery_threshold=1, latency_ms=1.0)
+            return reg.select()
+
+        pick = asyncio.run(promote())
+        assert pick is not None
+        assert pick.primary
+
+    def test_falls_back_to_secondary_when_primary_unhealthy(self):
+        reg = _registry()
+
+        async def run():
+            primary = reg.primaries()[0]
+            sec = next(r for r in reg.all() if not r.primary)
+            # Fail primary past threshold, succeed secondary.
+            for _ in range(3):
+                await reg.record_failure(
+                    primary, failure_threshold=3, reason="mock")
+            await reg.record_success(
+                sec, recovery_threshold=1, latency_ms=1.0)
+            return reg.select()
+
+        pick = asyncio.run(run())
+        assert pick is not None
+        assert not pick.primary
+
+    def test_failure_threshold_transitions_to_unhealthy(self):
+        reg = _registry()
+        primary = reg.primaries()[0]
+
+        async def run():
+            for _ in range(2):
+                await reg.record_failure(
+                    primary, failure_threshold=3, reason="x")
+            assert primary.health.status is not ReplicaStatus.UNHEALTHY
+            await reg.record_failure(
+                primary, failure_threshold=3, reason="x")
+            return primary
+
+        result = asyncio.run(run())
+        assert result.health.status is ReplicaStatus.UNHEALTHY
+        assert result.health.consecutive_failures == 3
+
+    def test_recovery_threshold_returns_to_healthy(self):
+        reg = _registry()
+        primary = reg.primaries()[0]
+
+        async def run():
+            for _ in range(3):
+                await reg.record_failure(
+                    primary, failure_threshold=3, reason="x")
+            assert primary.health.status is ReplicaStatus.UNHEALTHY
+            for _ in range(2):
+                await reg.record_success(
+                    primary, recovery_threshold=2, latency_ms=5.0)
+            return primary
+
+        result = asyncio.run(run())
+        assert result.health.status is ReplicaStatus.HEALTHY
+        assert result.health.last_latency_ms == 5.0
+
+    def test_record_success_resets_failure_counter(self):
+        reg = _registry()
+        primary = reg.primaries()[0]
+
+        async def run():
+            await reg.record_failure(
+                primary, failure_threshold=10, reason="x")
+            await reg.record_success(
+                primary, recovery_threshold=1, latency_ms=1.0)
+            return primary
+
+        result = asyncio.run(run())
+        assert result.health.consecutive_failures == 0
+
+
+class TestHealthCheckLoop:
+
+    def test_loop_transitions_replicas_on_probe_results(self):
+        config = RouterConfig(
+            replicas=[
+                ReplicaEndpoint(url="http://a", primary=True),
+                ReplicaEndpoint(url="http://b"),
+            ],
+            health_check_interval_seconds=0.01,
+            failure_threshold=1,
+            recovery_threshold=1,
+        )
+        reg = ReplicaRegistry(config.replicas)
+        stop_event = asyncio.Event()
+
+        calls: list[str] = []
+
+        async def probe(url, *, timeout):
+            calls.append(url)
+            if "http://a" in url:
+                return 1.0, None
+            return 0.0, "mock failure"
+
+        async def run() -> None:
+            task = asyncio.create_task(run_health_check_loop(
+                registry=reg, config=config, stop_event=stop_event,
+                http_get=probe,
+            ))
+            await asyncio.sleep(0.05)
+            stop_event.set()
+            await task
+
+        asyncio.run(run())
+        a = reg.get("http://a")
+        b = reg.get("http://b")
+        assert a is not None
+        assert b is not None
+        assert a.health.status is ReplicaStatus.HEALTHY
+        assert b.health.status is ReplicaStatus.UNHEALTHY
+        assert any("/health" in c for c in calls)
+
+
+class TestRouterApp:
+
+    def test_status_endpoint_lists_replicas(self):
+        from starlette.testclient import TestClient
+
+        from fastvideo.entrypoints.streaming.router.main import build_router_app
+
+        config = RouterConfig(
+            replicas=[
+                ReplicaEndpoint(url="http://a", primary=True),
+                ReplicaEndpoint(url="http://b"),
+            ],
+            health_check_interval_seconds=60,  # don't actually poll
+        )
+        reg = ReplicaRegistry(config.replicas)
+        app = build_router_app(config, registry=reg)
+        client = TestClient(app)
+        response = client.get("/status")
+        body = response.json()
+        urls = {r["url"] for r in body["replicas"]}
+        assert urls == {"http://a", "http://b"}
+        # Initial status is UNKNOWN.
+        assert all(r["status"] == "unknown" for r in body["replicas"])
+
+    def test_ws_rejects_when_no_healthy_replica(self):
+        from starlette.testclient import TestClient
+
+        from fastvideo.entrypoints.streaming.router.main import build_router_app
+
+        config = RouterConfig(
+            replicas=[ReplicaEndpoint(url="http://a", primary=True)],
+            health_check_interval_seconds=60,
+        )
+        reg = ReplicaRegistry(config.replicas)
+        app = build_router_app(config, registry=reg)
+        client = TestClient(app)
+        with client.websocket_connect("/v1/stream") as ws:
+            err = ws.receive_json()
+            assert err["type"] == "error"
+            assert err["code"] == "gpu_unavailable"
+
+
+class TestUnknownToHealthyImmediate:
+    """Initial probe must promote UNKNOWN -> HEALTHY without waiting for recovery_threshold."""
+
+    def test_first_success_promotes_unknown(self):
+        reg = _registry(num_primary=1, num_secondary=0)
+        primary = reg.primaries()[0]
+        assert primary.health.status is ReplicaStatus.UNKNOWN
+
+        async def run():
+            await reg.record_success(primary, recovery_threshold=10, latency_ms=1.0)
+
+        asyncio.run(run())
+        assert primary.health.status is ReplicaStatus.HEALTHY
+
+    def test_unhealthy_recovery_still_gated_by_threshold(self):
+        reg = _registry(num_primary=1, num_secondary=0)
+        primary = reg.primaries()[0]
+
+        async def run():
+            for _ in range(3):
+                await reg.record_failure(primary, failure_threshold=3, reason="x")
+            assert primary.health.status is ReplicaStatus.UNHEALTHY
+            await reg.record_success(primary, recovery_threshold=2, latency_ms=1.0)
+            assert primary.health.status is ReplicaStatus.UNHEALTHY  # 1/2 successes
+            await reg.record_success(primary, recovery_threshold=2, latency_ms=1.0)
+            assert primary.health.status is ReplicaStatus.HEALTHY    # 2/2 successes
+
+        asyncio.run(run())
+
+
+class TestConfigValidation:
+    """RouterConfig.__post_init__ rejects malformed configs."""
+
+    def test_rejects_path_in_url(self):
+        with pytest.raises(ValueError, match="without a path"):
+            RouterConfig(replicas=[ReplicaEndpoint(url="http://host:8000/api")])
+
+    def test_rejects_query_in_url(self):
+        with pytest.raises(ValueError, match="query/fragment"):
+            RouterConfig(replicas=[ReplicaEndpoint(url="http://host:8000?x=1")])
+
+    def test_rejects_fragment_in_url(self):
+        with pytest.raises(ValueError, match="query/fragment"):
+            RouterConfig(replicas=[ReplicaEndpoint(url="http://host:8000#frag")])
+
+    def test_rejects_duplicate_urls(self):
+        with pytest.raises(ValueError, match="Duplicate"):
+            RouterConfig(replicas=[
+                ReplicaEndpoint(url="http://host:8000"),
+                ReplicaEndpoint(url="http://host:8000"),
+            ])
+
+    def test_accepts_trailing_slash(self):
+        # parsed.path == "/" should be allowed
+        cfg = RouterConfig(replicas=[ReplicaEndpoint(url="http://host:8000/")])
+        assert cfg.replicas[0].url == "http://host:8000/"

pyproject.toml

@@ -145,6 +145,7 @@ prompt-enhancer = [
 streaming = [
     "fastvideo[prompt-enhancer]",
     "fastvideo[prompt-safety]",
+    "websockets",
 ]
 
 rocm = [