websocket-proxy 0.1.1__py3-none-any.whl

websocket_proxy/__init__.py

@@ -0,0 +1,13 @@
+"""Fox Bridge - Foxglove WebSocket Proxy with Message Transformations"""
+
+from websocket_proxy.proxy import ProxyBridge
+from websocket_proxy.transformers import TransformerRegistry
+from websocket_proxy.transformers.image_to_video import ImageToVideoTransformer
+from websocket_proxy.transformers.pointcloud_voxel import PointCloudVoxelTransformer
+
+__all__ = [
+    "ImageToVideoTransformer",
+    "PointCloudVoxelTransformer",
+    "ProxyBridge",
+    "TransformerRegistry",
+]

websocket_proxy/__main__.py

@@ -0,0 +1,228 @@
+import argparse
+import asyncio
+import contextlib
+import logging
+import signal
+import sys
+
+from rich.console import Console
+from rich.logging import RichHandler
+
+from websocket_proxy.dashboard import DashboardRenderer
+from websocket_proxy.proxy import ProxyBridge
+from websocket_proxy.transformers import TransformerRegistry
+from websocket_proxy.transformers.image_to_video import ImageToVideoTransformer
+from websocket_proxy.transformers.pointcloud_voxel import PointCloudVoxelTransformer
+
+logger = logging.getLogger(__name__)
+
+
+def parse_args() -> argparse.Namespace:
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser(
+        description="Foxglove WebSocket proxy - forwards topics with optional transformations"
+    )
+    parser.add_argument(
+        "source_ws",
+        help="WebSocket URL of the upstream Foxglove bridge (e.g., ws://localhost:8765)",
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=8766,
+        help="Port to listen on for downstream clients (default: 8766)",
+    )
+    parser.add_argument(
+        "--host",
+        default="0.0.0.0",  # noqa: S104
+        help="Host to listen on for downstream clients (default: 0.0.0.0)",
+    )
+    parser.add_argument(
+        "--verbose",
+        "-v",
+        action="store_true",
+        help="Enable verbose debug logging",
+    )
+    parser.add_argument(
+        "--throttle-hz",
+        type=float,
+        default=1.0,
+        help="Topic throttle rate in Hz (default: 1.0; set to 0 to disable)",
+    )
+    parser.add_argument(
+        "--max-message-size",
+        type=int,
+        default=0,
+        help="Maximum websocket message size in bytes (<=0 disables limit, default: unlimited)",
+    )
+
+    parser.add_argument(
+        "--image-codec",
+        default="h264",
+        help="Video codec to use for image compression (default: h264)",
+    )
+    parser.add_argument(
+        "--image-quality",
+        type=int,
+        default=23,
+        help="CRF/quality value for image compression (lower is higher quality, default: 23)",
+    )
+    parser.add_argument(
+        "--image-preset",
+        default="fast",
+        help="Encoder preset for image compression (default: fast)",
+    )
+    parser.add_argument(
+        "--image-max-dimension",
+        type=int,
+        default=480,
+        help="Maximum width/height used when downscaling images before encoding (default: 480)",
+    )
+    parser.add_argument(
+        "--image-disable-hw",
+        dest="image_use_hardware",
+        action="store_false",
+        help="Disable hardware acceleration for image compression",
+    )
+
+    parser.add_argument(
+        "--pointcloud-voxel-size",
+        type=float,
+        default=0.1,
+        help="Voxel size (in meters) for point cloud downsampling (default: 0.1)",
+    )
+    parser.add_argument(
+        "--pointcloud-keep-nans",
+        dest="pointcloud_skip_nans",
+        action="store_false",
+        help="Keep NaN points when voxelizing point clouds (default: drop NaNs)",
+    )
+
+    parser.add_argument(
+        "--no-dashboard",
+        action="store_true",
+        help="Disable the live dashboard display",
+    )
+    parser.add_argument(
+        "--dashboard-refresh-rate",
+        type=float,
+        default=1.0,
+        help="Dashboard refresh rate in seconds (default: 1.0)",
+    )
+
+    parser.set_defaults(image_use_hardware=True, pointcloud_skip_nans=True)
+    return parser.parse_args()
+
+
+async def main_async(args: argparse.Namespace) -> None:
+    """Async main function."""
+    # Create shared console for dashboard and logging
+    console = Console()
+
+    # Create transformer registry and register transformers (BEFORE configuring logging)
+    registry = TransformerRegistry()
+
+    # Register image to video transformer
+    image_transformer = ImageToVideoTransformer(
+        codec=args.image_codec,
+        quality=args.image_quality,
+        preset=args.image_preset,
+        use_hardware=args.image_use_hardware,
+        max_dimension=args.image_max_dimension,
+    )
+    registry.register(image_transformer)
+
+    pointcloud_transformer = PointCloudVoxelTransformer(
+        voxel_size=args.pointcloud_voxel_size,
+        skip_nans=args.pointcloud_skip_nans,
+    )
+    registry.register(pointcloud_transformer)
+
+    # Create proxy bridge with transformers
+    bridge = ProxyBridge(
+        upstream_url=args.source_ws,
+        listen_host=args.host,
+        listen_port=args.port,
+        transformer_registry=registry,
+        default_throttle_hz=args.throttle_hz,
+        max_message_size=args.max_message_size if args.max_message_size > 0 else None,
+    )
+
+    # Create dashboard if enabled (with shared console for logging integration)
+    dashboard = None
+    if not args.no_dashboard:
+        dashboard = DashboardRenderer(
+            bridge, refresh_rate=args.dashboard_refresh_rate, console=console
+        )
+        # Start dashboard BEFORE configuring logging
+        dashboard.start_sync()
+
+    # NOW configure logging with Rich handler (after dashboard is started)
+    logging.basicConfig(
+        level=logging.DEBUG if args.verbose else logging.INFO,
+        format="%(message)s",
+        datefmt="[%X]",
+        handlers=[
+            RichHandler(
+                console=console,
+                rich_tracebacks=True,
+                tracebacks_show_locals=args.verbose,
+            )
+        ],
+    )
+
+    logger.info("Registered transformers:")
+    for transformer in registry.get_all_transformers():
+        logger.info(f"  {transformer.get_input_schema()} -> {transformer.get_output_schema()}")
+
+    # Setup signal handlers for graceful shutdown
+    loop = asyncio.get_running_loop()
+
+    def signal_handler() -> None:
+        logger.info("Received shutdown signal")
+        asyncio.create_task(bridge.stop())  # noqa: RUF006
+
+    for sig in (signal.SIGTERM, signal.SIGINT):
+        loop.add_signal_handler(sig, signal_handler)
+
+    try:
+        if dashboard:
+            # Dashboard already started above (before logging config)
+            # Just create a background task for dashboard updates
+            dashboard_task = asyncio.create_task(dashboard.run_updates())
+
+            try:
+                # Start proxy (this will block until stop() is called)
+                await bridge.start()
+            finally:
+                # Cancel dashboard updates
+                dashboard_task.cancel()
+                with contextlib.suppress(asyncio.CancelledError):
+                    await dashboard_task
+                await dashboard.stop()
+        else:
+            # No dashboard - just start proxy
+            await bridge.start()
+    except KeyboardInterrupt:
+        logger.info("Keyboard interrupt received")
+    except Exception:
+        logger.exception("Unexpected error in proxy bridge")
+        sys.exit(1)
+    finally:
+        await bridge.stop()
+        if dashboard:
+            await dashboard.stop()
+
+
+def main() -> None:
+    """Main entry point."""
+    args = parse_args()
+
+    try:
+        asyncio.run(main_async(args))
+    except KeyboardInterrupt:
+        logger.info("Exiting")
+
+
+if __name__ == "__main__":
+    main()

websocket_proxy/dashboard.py

@@ -0,0 +1,245 @@
+"""Rich-based dashboard for websocket proxy server."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from datetime import datetime, timedelta, timezone
+from typing import TYPE_CHECKING
+
+from rich.console import Console, Group
+from rich.live import Live
+from rich.panel import Panel
+from rich.table import Table
+from rich.text import Text
+
+if TYPE_CHECKING:
+    from .metrics import MetricsCollector
+    from .proxy import ProxyBridge
+
+logger = logging.getLogger(__name__)
+
+
+def _format_duration(seconds: float) -> str:
+    """Format duration in seconds to human-readable string."""
+    if seconds < 60:
+        return f"{seconds:.0f}s"
+    if seconds < 3600:
+        return f"{seconds // 60:.0f}m {seconds % 60:.0f}s"
+    hours = seconds // 3600
+    minutes = (seconds % 3600) // 60
+    return f"{hours:.0f}h {minutes:.0f}m"
+
+
+def _format_rate(rate: float) -> str:
+    """Format rate (msgs/sec) to human-readable string."""
+    if rate < 0.01:
+        return "0.00"
+    if rate < 1:
+        return f"{rate:.2f}"
+    if rate < 10:
+        return f"{rate:.1f}"
+    return f"{rate:.0f}"
+
+
+def _format_bandwidth(bytes_per_sec: float) -> str:
+    """Format bandwidth to human-readable string."""
+    if bytes_per_sec < 1024:
+        return f"{bytes_per_sec:.0f} B/s"
+    if bytes_per_sec < 1024 * 1024:
+        return f"{bytes_per_sec / 1024:.1f} KB/s"
+    return f"{bytes_per_sec / (1024 * 1024):.2f} MB/s"
+
+
+def _format_bytes(byte_count: int) -> str:
+    """Format byte count to human-readable string."""
+    if byte_count < 1024:
+        return f"{byte_count} B"
+    if byte_count < 1024 * 1024:
+        return f"{byte_count / 1024:.1f} KB"
+    if byte_count < 1024 * 1024 * 1024:
+        return f"{byte_count / (1024 * 1024):.1f} MB"
+    return f"{byte_count / (1024 * 1024 * 1024):.2f} GB"
+
+
+def _format_timestamp(dt: datetime | None) -> str:
+    """Format timestamp to relative time string."""
+    if dt is None:
+        return "never"
+
+    now = datetime.now(timezone.utc)
+    diff = now - dt
+
+    if diff < timedelta(seconds=1):
+        return "just now"
+    if diff < timedelta(seconds=60):
+        return f"{diff.seconds}s ago"
+    if diff < timedelta(minutes=60):
+        return f"{diff.seconds // 60}m ago"
+    return dt.strftime("%H:%M:%S")
+
+
+class DashboardRenderer:
+    """Renders a live dashboard for the proxy server using Rich."""
+
+    def __init__(self, proxy: ProxyBridge, refresh_rate: float, console: Console) -> None:
+        self.proxy = proxy
+        self.metrics: MetricsCollector = proxy.metrics
+        self.refresh_rate = refresh_rate
+        self.console = console
+        self._live: Live | None = None
+
+    def _create_header_panel(self) -> Panel:
+        """Create the header panel with global stats."""
+        uptime = _format_duration(self.metrics.get_uptime())
+        total_clients = len(self.metrics.clients)
+
+        header_text = Text()
+        header_text.append("Foxglove WebSocket Proxy Dashboard", style="bold cyan")
+        header_text.append("\n\n")
+        header_text.append("Uptime: ", style="bold")
+        header_text.append(uptime)
+        header_text.append("  |  Connected Clients: ", style="bold")
+        header_text.append(str(total_clients), style="green" if total_clients > 0 else "dim")
+
+        # Add upstream metrics
+        header_text.append("\n\n")
+        header_text.append("Upstream Status: ", style="bold")
+        if self.metrics.upstream_connected:
+            header_text.append("● Connected", style="green bold")
+        else:
+            header_text.append("● Disconnected", style="red bold")
+
+        header_text.append("  |  Topics: ", style="bold")
+        header_text.append(str(self.metrics.upstream_topic_count), style="cyan")
+
+        header_text.append("  |  Transformed: ", style="bold")
+        header_text.append(str(self.metrics.transformed_channel_count), style="magenta")
+
+        header_text.append("\n")
+        header_text.append("Messages Received: ", style="bold")
+        header_text.append(str(self.metrics.upstream_messages_received), style="green")
+
+        header_text.append("  |  Throttled: ", style="bold")
+        throttled_style = "red" if self.metrics.upstream_messages_throttled > 0 else "dim"
+        header_text.append(str(self.metrics.upstream_messages_throttled), style=throttled_style)
+
+        return Panel(header_text, border_style="blue")
+
+    def _create_clients_table(self) -> Table:
+        """Create the clients table."""
+        table = Table(
+            title="Connected Clients",
+            title_style="bold magenta",
+            show_header=True,
+            header_style="bold",
+            show_lines=False,
+            expand=False,
+        )
+
+        table.add_column("ID", style="cyan", no_wrap=True)
+        table.add_column("Remote Address", style="blue", no_wrap=True)
+        table.add_column("Connected", style="green", no_wrap=True)
+        table.add_column("Msg/s", justify="right", style="yellow")
+        table.add_column("Bandwidth", justify="right", style="yellow")
+        table.add_column("Msgs", justify="right", style="white")
+        table.add_column("Bytes", justify="right", style="yellow")
+        table.add_column("Subs", justify="center", style="cyan")
+        table.add_column("Errors", justify="center", style="red")
+        table.add_column("Last Msg", style="dim", no_wrap=True)
+
+        # Sort clients by connection time (oldest first)
+        sorted_clients = sorted(
+            self.metrics.clients.values(),
+            key=lambda c: c.connected_at,
+        )
+
+        for client in sorted_clients:
+            client_id_short = client.client_id.replace("client_", "")[:12]
+            duration = _format_duration(client.connected_duration)
+            msg_rate = _format_rate(client.get_message_rate())
+            bandwidth = _format_bandwidth(client.get_bandwidth())
+            bytes_send = _format_bytes(client.bytes_sent)
+            last_msg = _format_timestamp(client.last_message_at)
+
+            # Style errors in red if > 0
+            errors_str = str(client.errors)
+            errors_style = "red bold" if client.errors > 0 else "dim"
+
+            table.add_row(
+                client_id_short,
+                client.remote_address,
+                duration,
+                msg_rate,
+                bandwidth,
+                str(client.messages_sent),
+                bytes_send,
+                str(client.subscription_count),
+                Text(errors_str, style=errors_style),
+                last_msg,
+            )
+
+        if not sorted_clients:
+            table.add_row(
+                Text("No clients connected", style="dim italic"),
+                "",
+                "",
+                "",
+                "",
+                "",
+                "",
+                "",
+                "",
+                "",
+            )
+
+        return table
+
+    def _create_layout(self) -> Panel:
+        """Create the full dashboard layout."""
+        header = self._create_header_panel()
+        clients_table = self._create_clients_table()
+
+        # Group everything together
+        layout = Group(
+            header,
+            "",  # Spacer
+            clients_table,
+        )
+
+        return Panel(layout, border_style="bright_blue", padding=(1, 2))
+
+    def start_sync(self) -> None:
+        """Start the live dashboard (synchronous - starts the Live display)."""
+        self._live = Live(
+            self._create_layout(),
+            console=self.console,
+            refresh_per_second=1 / self.refresh_rate,
+            screen=False,
+            auto_refresh=True,  # Enable auto-refresh so updates are visible
+        )
+        self._live.start()
+        # Do an initial render
+        self._live.update(self._create_layout())
+
+    async def run_updates(self) -> None:
+        """Run the dashboard update loop (call after start_sync)."""
+        if self._live is None:
+            return
+
+        while True:
+            try:
+                self._live.update(self._create_layout())
+                await asyncio.sleep(self.refresh_rate)
+            except asyncio.CancelledError:
+                break
+            except (RuntimeError, ValueError, KeyError, IndexError) as e:
+                # Log errors during rendering but continue
+                logger.debug("Dashboard rendering error: %s", e)
+                await asyncio.sleep(self.refresh_rate)
+
+    async def stop(self) -> None:
+        """Stop the live dashboard."""
+        if self._live:
+            self._live.stop()
+            self._live = None

websocket_proxy/metrics.py

@@ -0,0 +1,119 @@
+"""Metrics collection for websocket proxy server."""
+
+import time
+from collections import deque
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+
+
+@dataclass
+class ClientMetrics:
+    """Metrics for a single client connection."""
+
+    client_id: str
+    remote_address: str
+    connected_at: datetime
+    user_agent: str | None = None
+
+    # Traffic metrics
+    messages_sent: int = 0
+    bytes_sent: int = 0
+    errors: int = 0
+    last_message_at: datetime | None = None
+
+    # Subscriptions
+    subscription_count: int = 0
+    subscribed_topics: set[str] = field(default_factory=set)
+
+    # Rate tracking (samples stored for windowed calculations)
+    _message_samples: deque[tuple[float, int]] = field(default_factory=lambda: deque(maxlen=60))
+    _byte_samples: deque[tuple[float, int]] = field(default_factory=lambda: deque(maxlen=60))
+
+    def record_message(self, byte_count: int) -> None:
+        """Record a message sent to this client."""
+        now = time.time()
+        self.messages_sent += 1
+        self.bytes_sent += byte_count
+        self.last_message_at = datetime.now(timezone.utc)
+        self._message_samples.append((now, 1))
+        self._byte_samples.append((now, byte_count))
+
+    def record_error(self) -> None:
+        """Record an error for this client."""
+        self.errors += 1
+
+    def get_message_rate(self, window_seconds: float = 5.0) -> float:
+        """Calculate messages per second over the last N seconds."""
+        return self._calculate_rate(self._message_samples, window_seconds)
+
+    def get_bandwidth(self, window_seconds: float = 5.0) -> float:
+        """Calculate bytes per second over the last N seconds."""
+        return self._calculate_rate(self._byte_samples, window_seconds)
+
+    @staticmethod
+    def _calculate_rate(samples: deque[tuple[float, int]], window_seconds: float) -> float:
+        """Calculate rate from time-stamped samples within window."""
+        if not samples:
+            return 0.0
+
+        now = time.time()
+        cutoff = now - window_seconds
+        total = sum(count for ts, count in samples if ts >= cutoff)
+
+        # Find actual time span of samples in window
+        valid_samples = [(ts, count) for ts, count in samples if ts >= cutoff]
+        if not valid_samples:
+            return 0.0
+
+        oldest_ts = valid_samples[0][0]
+        time_span = now - oldest_ts
+
+        if time_span <= 0:
+            return 0.0
+
+        return total / time_span
+
+    @property
+    def connected_duration(self) -> float:
+        """Get connection duration in seconds."""
+        return (datetime.now(timezone.utc) - self.connected_at).total_seconds()
+
+
+class MetricsCollector:
+    """Central metrics collector for the proxy server."""
+
+    def __init__(self) -> None:
+        self.clients: dict[str, ClientMetrics] = {}
+        self._start_time = datetime.now(timezone.utc)
+
+        # Upstream metrics
+        self.upstream_connected: bool = False
+        self.upstream_topic_count: int = 0
+        self.upstream_messages_received: int = 0
+        self.upstream_messages_throttled: int = 0
+        self.transformed_channel_count: int = 0
+
+    def add_client(
+        self, client_id: str, remote_address: str, user_agent: str | None = None
+    ) -> ClientMetrics:
+        """Add a new client to track."""
+        metrics = ClientMetrics(
+            client_id=client_id,
+            remote_address=remote_address,
+            connected_at=datetime.now(timezone.utc),
+            user_agent=user_agent,
+        )
+        self.clients[client_id] = metrics
+        return metrics
+
+    def remove_client(self, client_id: str) -> None:
+        """Remove a client from tracking."""
+        self.clients.pop(client_id, None)
+
+    def get_client(self, client_id: str) -> ClientMetrics | None:
+        """Get client metrics by ID."""
+        return self.clients.get(client_id)
+
+    def get_uptime(self) -> float:
+        """Get server uptime in seconds."""
+        return (datetime.now(timezone.utc) - self._start_time).total_seconds()