"""Thin IQ-streaming agent. Listens for control messages from the RIA Hub over a persistent WebSocket. When the server sends ``start``, opens the SDR described in ``radio_config``, loops over ``sdr.rx(buffer_size)``, and sends each buffer as raw interleaved float32 bytes. ``stop`` closes the SDR; ``configure`` applies parameter updates at the next capture boundary. """ from __future__ import annotations import asyncio import logging from typing import Any import numpy as np from .hardware import heartbeat_payload from .ws_client import WsClient logger = logging.getLogger("ria_agent.streamer") _DEFAULT_BUFFER_SIZE = 1024 class Streamer: """Main streamer loop. Parameters ---------- ws: Connected :class:`WsClient`. sdr_factory: Callable ``(device, identifier) -> SDR``. Defaults to :func:`ria_toolkit_oss.sdr.get_sdr_device`. Injectable for tests. """ def __init__(self, ws: WsClient, sdr_factory=None) -> None: self.ws = ws self._sdr_factory = sdr_factory self._app_id: str | None = None self._sdr: Any = None self._pending_config: dict = {} self._capture_task: asyncio.Task | None = None self._status = "idle" # ------------------------------------------------------------------ # WsClient wiring def build_heartbeat(self) -> dict: return heartbeat_payload(status=self._status, app_id=self._app_id) async def on_message(self, msg: dict) -> None: t = msg.get("type") if t == "start": await self._handle_start(msg) elif t == "stop": await self._handle_stop(msg) elif t == "configure": self._pending_config.update(msg.get("radio_config") or {}) logger.debug("Queued configure: %s", self._pending_config) else: logger.warning("Unknown server message type: %r", t) # ------------------------------------------------------------------ async def _handle_start(self, msg: dict) -> None: if self._capture_task is not None and not self._capture_task.done(): logger.warning("start received while already streaming — ignoring") return self._app_id = msg.get("app_id") radio_config = dict(msg.get("radio_config") or {}) device = radio_config.pop("device", None) identifier = radio_config.pop("identifier", None) buffer_size = int(radio_config.pop("buffer_size", _DEFAULT_BUFFER_SIZE)) if not device: await self._send_error("start missing radio_config.device") return try: factory = self._sdr_factory or _default_sdr_factory self._sdr = factory(device, identifier) _apply_sdr_config(self._sdr, radio_config) except Exception as exc: logger.exception("Failed to open SDR %r", device) await self._send_error(f"SDR init failed: {exc}") return self._status = "streaming" await self._send_status("streaming") self._capture_task = asyncio.create_task( self._capture_loop(buffer_size), name="ria-streamer-capture" ) async def _handle_stop(self, msg: dict) -> None: if self._capture_task is not None: self._capture_task.cancel() try: await self._capture_task except (asyncio.CancelledError, Exception): pass self._capture_task = None self._close_sdr() self._app_id = None self._status = "idle" await self._send_status("idle") async def _capture_loop(self, buffer_size: int) -> None: loop = asyncio.get_running_loop() try: while True: if self._pending_config: cfg = self._pending_config self._pending_config = {} try: _apply_sdr_config(self._sdr, cfg) except Exception as exc: logger.warning("Applying configure failed: %s", exc) try: samples = await loop.run_in_executor(None, self._sdr.rx, buffer_size) except Exception as exc: from ria_toolkit_oss.sdr import SdrDisconnectedError if isinstance(exc, SdrDisconnectedError): logger.warning("SDR disconnected: %s", exc) await self._send_error(f"SDR disconnected: {exc}") else: logger.exception("SDR rx error") await self._send_error(f"SDR capture failed: {exc}") break payload = _samples_to_interleaved_float32(samples) try: await self.ws.send_bytes(payload) except Exception as exc: logger.warning("Send failed: %s — ending capture", exc) break except asyncio.CancelledError: raise finally: self._close_sdr() def _close_sdr(self) -> None: if self._sdr is None: return try: self._sdr.close() except Exception: pass self._sdr = None async def _send_status(self, status: str) -> None: try: await self.ws.send_json({"type": "status", "status": status, "app_id": self._app_id}) except Exception as exc: logger.debug("Status send failed: %s", exc) async def _send_error(self, message: str) -> None: try: await self.ws.send_json({"type": "error", "app_id": self._app_id, "message": message}) except Exception as exc: logger.debug("Error-frame send failed: %s", exc) # --------------------------------------------------------------------------- # Helpers _CONFIG_ATTR_MAP = { "sample_rate": ("sample_rate", "rx_sample_rate"), "center_frequency": ("center_freq", "rx_center_frequency"), "center_freq": ("center_freq", "rx_center_frequency"), "gain": ("gain", "rx_gain"), "bandwidth": ("bandwidth", "rx_bandwidth"), } def _apply_sdr_config(sdr: Any, cfg: dict) -> None: """Apply a radio_config dict to an SDR, trying multiple attribute aliases.""" for key, value in cfg.items(): if value is None: continue attrs = _CONFIG_ATTR_MAP.get(key, (key,)) applied = False for attr in attrs: if hasattr(sdr, attr): try: setattr(sdr, attr, value) applied = True break except Exception as exc: logger.debug("setattr %s=%r failed: %s", attr, value, exc) if not applied: logger.debug("radio_config key %r ignored (no matching attr)", key) def _samples_to_interleaved_float32(samples: Any) -> bytes: """Convert complex IQ samples (any numeric dtype) to interleaved float32 bytes.""" arr = np.asarray(samples) if np.iscomplexobj(arr): interleaved = np.empty(arr.size * 2, dtype=np.float32) interleaved[0::2] = arr.real.astype(np.float32, copy=False).ravel() interleaved[1::2] = arr.imag.astype(np.float32, copy=False).ravel() return interleaved.tobytes() return arr.astype(np.float32, copy=False).tobytes() def _default_sdr_factory(device: str, identifier: str | None): from ria_toolkit_oss.sdr import get_sdr_device return get_sdr_device(device, ident=identifier) # --------------------------------------------------------------------------- # Top-level entry async def run_streamer(ws_url: str, token: str) -> None: """Connect to *ws_url* and run the streamer loop until cancelled.""" ws = WsClient(ws_url, token) streamer = Streamer(ws) await ws.run(streamer.on_message, streamer.build_heartbeat)