"""IQ-streaming agent. Listens for control messages from the RIA Hub over a persistent WebSocket. Supports: - An **RX session** (hub sends ``start``/``stop``/``configure``; agent opens the SDR, loops ``sdr.rx()`` and ships raw interleaved float32 IQ). - A **TX session** (hub sends ``tx_start``/``tx_stop``/``tx_configure`` plus binary IQ frames; agent feeds them into ``sdr._stream_tx``). Phase 3 wires up the session plumbing and rejects TX when ``cfg.tx_enabled`` is False; Phase 4 implements the full TX loop. Both sessions can run concurrently on the same physical SDR (FDD) — a ref-counted SDR registry shares one driver instance when RX and TX name the same ``(device, identifier)``. """ from __future__ import annotations import asyncio import logging import queue import threading import time from dataclasses import dataclass, field from typing import Any import numpy as np from .config import AgentConfig from .hardware import heartbeat_payload from .ws_client import WsClient logger = logging.getLogger("ria_agent.streamer") _DEFAULT_BUFFER_SIZE = 1024 # --------------------------------------------------------------------------- # Session dataclasses @dataclass class RxSession: app_id: str sdr: Any device_key: tuple[str, str | None] buffer_size: int task: asyncio.Task | None = None pending_config: dict = field(default_factory=dict) @dataclass class TxSession: app_id: str sdr: Any device_key: tuple[str, str | None] buffer_size: int task: Any = None # concurrent.futures.Future from run_in_executor pending_config: dict = field(default_factory=dict) underrun_policy: str = "pause" last_buffer: np.ndarray | None = None stop_event: threading.Event = field(default_factory=threading.Event) started_at: float = 0.0 max_duration_s: float | None = None state: str = "armed" # Thread-safe queue of inbound interleaved-float32 IQ frames. Bounded so # hub-side over-production triggers WS backpressure rather than memory # growth in the agent. in_queue: "queue.Queue[bytes]" = field(default_factory=lambda: queue.Queue(maxsize=8)) # Set by the TX callback when it hits an underrun while policy=="pause"; # asyncio side flips the session state and emits tx_status. underrun_flag: threading.Event = field(default_factory=threading.Event) # --------------------------------------------------------------------------- # SDR registry (ref-counted so one Pluto handle serves RX + TX simultaneously) class _SdrRegistry: def __init__(self, factory): self._factory = factory self._instances: dict[tuple[str, str | None], tuple[Any, int]] = {} self._lock = threading.Lock() def acquire(self, device: str, identifier: str | None) -> tuple[Any, tuple[str, str | None]]: key = (device, identifier) with self._lock: if key in self._instances: sdr, rc = self._instances[key] self._instances[key] = (sdr, rc + 1) return sdr, key # Build outside the lock: driver init can be slow and we don't want to # block concurrent releases on unrelated devices. sdr = self._factory(device, identifier) with self._lock: if key in self._instances: # Raced another acquirer; discard our duplicate and share theirs. other_sdr, rc = self._instances[key] try: sdr.close() except Exception: pass self._instances[key] = (other_sdr, rc + 1) return other_sdr, key self._instances[key] = (sdr, 1) return sdr, key def release(self, key: tuple[str, str | None]) -> bool: """Decrement refcount. Returns True if the caller owns the last reference and should close the SDR.""" with self._lock: sdr, rc = self._instances.get(key, (None, 0)) if sdr is None: return False if rc <= 1: del self._instances[key] return True self._instances[key] = (sdr, rc - 1) return False def refcount(self, key: tuple[str, str | None]) -> int: with self._lock: return self._instances.get(key, (None, 0))[1] # --------------------------------------------------------------------------- # Streamer class Streamer: """Main streamer loop. Parameters ---------- ws: Connected :class:`WsClient`. sdr_factory: Callable ``(device, identifier) -> SDR``. Defaults to the helper in :mod:`ria_toolkit_oss.sdr`. Injectable for tests. cfg: :class:`AgentConfig` for interlocks (``tx_enabled`` and caps) and heartbeat capabilities. Defaults to an empty ``AgentConfig()`` which leaves TX disabled. """ def __init__( self, ws, sdr_factory=None, cfg: AgentConfig | None = None, ) -> None: self.ws = ws self._cfg = cfg or AgentConfig() self._registry = _SdrRegistry(sdr_factory or _default_sdr_factory) self._rx: RxSession | None = None self._tx: TxSession | None = None # Pending radio_config accepted via ``configure`` before ``start``. self._standalone_pending_config: dict = {} # Cached asyncio event loop, set the first time a handler runs. Used # to schedule async callbacks from the TX executor thread. self._loop: asyncio.AbstractEventLoop | None = None # ------------------------------------------------------------------ # Back-compat read-only shims for callers that check ``._sdr`` etc. # Writes to these attributes are not supported — use the session objects. @property def _sdr(self): return self._rx.sdr if self._rx is not None else None @property def _pending_config(self) -> dict: return self._rx.pending_config if self._rx is not None else self._standalone_pending_config # ------------------------------------------------------------------ # WsClient wiring def build_heartbeat(self) -> dict: status = "streaming" if (self._rx is not None or self._tx is not None) else "idle" app_id: str | None = None if self._rx is not None: app_id = self._rx.app_id elif self._tx is not None: app_id = self._tx.app_id sessions: dict[str, dict] = {} if self._rx is not None: sessions["rx"] = {"app_id": self._rx.app_id, "state": "streaming"} if self._tx is not None: sessions["tx"] = {"app_id": self._tx.app_id, "state": self._tx.state} return heartbeat_payload( status=status, app_id=app_id, cfg=self._cfg, sessions=sessions or None, ) async def on_message(self, msg: dict) -> None: t = msg.get("type") handler = { "start": self._handle_rx_start, "stop": self._handle_rx_stop, "configure": self._handle_rx_configure, "tx_start": self._handle_tx_start, "tx_stop": self._handle_tx_stop, "tx_configure": self._handle_tx_configure, }.get(t) if handler is None: logger.warning("Unknown server message type: %r", t) return await handler(msg) async def on_binary(self, data: bytes) -> None: tx = self._tx if tx is None: logger.debug("Dropping %d-byte binary frame: no TX session", len(data)) return # Backpressure: if the TX queue is full, await briefly so the hub's # ``await ws.send`` throttles naturally via TCP. We don't block # indefinitely — a 2s stall means something else is wrong. loop = asyncio.get_running_loop() try: await loop.run_in_executor(None, lambda: tx.in_queue.put(data, timeout=2.0)) except queue.Full: logger.warning("TX queue stalled; dropping frame") # ================================================================== # RX async def _handle_rx_start(self, msg: dict) -> None: if self._rx is not None: logger.warning("start received while already streaming — ignoring") return app_id = msg.get("app_id") or "" 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(app_id, "start missing radio_config.device") return try: sdr, device_key = self._registry.acquire(device, identifier) _apply_sdr_config(sdr, radio_config) except Exception as exc: logger.exception("Failed to open SDR %r", device) await self._send_error(app_id, f"SDR init failed: {exc}") return # Inherit any pending config that was queued before start. pending = dict(self._standalone_pending_config) self._standalone_pending_config = {} session = RxSession( app_id=app_id, sdr=sdr, device_key=device_key, buffer_size=buffer_size, pending_config=pending, ) self._rx = session await self._send_status("streaming", app_id) session.task = asyncio.create_task( self._capture_loop(session), name="ria-streamer-capture" ) async def _handle_rx_stop(self, msg: dict) -> None: session = self._rx if session is None: return if session.task is not None: session.task.cancel() try: await session.task except (asyncio.CancelledError, Exception): pass self._close_session_sdr(session) app_id = session.app_id self._rx = None await self._send_status("idle", app_id) async def _handle_rx_configure(self, msg: dict) -> None: cfg = dict(msg.get("radio_config") or {}) if self._rx is not None: self._rx.pending_config.update(cfg) else: self._standalone_pending_config.update(cfg) logger.debug("Queued configure: %s", cfg) async def _capture_loop(self, session: RxSession) -> None: loop = asyncio.get_running_loop() try: while True: if session.pending_config: cfg = session.pending_config session.pending_config = {} try: _apply_sdr_config(session.sdr, cfg) except Exception as exc: logger.warning("Applying configure failed: %s", exc) try: samples = await loop.run_in_executor( None, session.sdr.rx, session.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(session.app_id, f"SDR disconnected: {exc}") else: logger.exception("SDR rx error") await self._send_error(session.app_id, 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_session_sdr(session) # If the loop died on its own (e.g. SDR disconnect), clear the # session handle so future ``start`` messages can proceed. if self._rx is session: self._rx = None # ================================================================== # TX async def _handle_tx_start(self, msg: dict) -> None: app_id = msg.get("app_id") or "" radio_config = dict(msg.get("radio_config") or {}) # --- interlocks (agent-enforced; never trust the hub alone) --- if not self._cfg.tx_enabled: await self._send_tx_status(app_id, "error", "tx disabled on this agent") return tx_gain = radio_config.get("tx_gain") if ( self._cfg.tx_max_gain_db is not None and tx_gain is not None and float(tx_gain) > float(self._cfg.tx_max_gain_db) ): await self._send_tx_status( app_id, "error", f"tx_gain {tx_gain} exceeds cap {self._cfg.tx_max_gain_db}", ) return tx_freq = radio_config.get("tx_center_frequency") if self._cfg.tx_allowed_freq_ranges and tx_freq is not None: f = float(tx_freq) if not any(float(lo) <= f <= float(hi) for lo, hi in self._cfg.tx_allowed_freq_ranges): await self._send_tx_status( app_id, "error", f"tx_center_frequency {tx_freq} outside allowed ranges", ) return if self._tx is not None: await self._send_tx_status(app_id, "error", "tx already active on this agent") return # --- device --- device = radio_config.pop("device", None) identifier = radio_config.pop("identifier", None) buffer_size = int(radio_config.pop("buffer_size", _DEFAULT_BUFFER_SIZE)) underrun_policy = str(radio_config.pop("underrun_policy", "pause")) if underrun_policy not in ("pause", "zero", "repeat"): await self._send_tx_status( app_id, "error", f"invalid underrun_policy {underrun_policy!r}" ) return if not device: await self._send_tx_status(app_id, "error", "tx_start missing radio_config.device") return device_key: tuple[str, str | None] | None = None sdr: Any = None try: sdr, device_key = self._registry.acquire(device, identifier) _apply_sdr_config(sdr, radio_config) # Only call init_tx when the hub supplied the three required # parameters. Drivers that gate _stream_tx on _tx_initialized # (e.g. Pluto) need this; drivers that don't (e.g. Mock) tolerate # its absence. init_args = { k: radio_config.get(f"tx_{k}") for k in ("sample_rate", "center_frequency", "gain") } if hasattr(sdr, "init_tx") and all(v is not None for v in init_args.values()): sdr.init_tx( sample_rate=init_args["sample_rate"], center_frequency=init_args["center_frequency"], gain=init_args["gain"], channel=radio_config.get("tx_channel", 0), gain_mode=radio_config.get("tx_gain_mode", "manual"), ) except Exception as exc: if device_key is not None: if self._registry.release(device_key): try: sdr.close() except Exception: pass logger.exception("Failed to init TX on %r", device) await self._send_tx_status(app_id, "error", f"tx init failed: {exc}") return self._loop = asyncio.get_running_loop() session = TxSession( app_id=app_id, sdr=sdr, device_key=device_key, buffer_size=buffer_size, underrun_policy=underrun_policy, started_at=time.monotonic(), max_duration_s=self._cfg.tx_max_duration_s, ) self._tx = session await self._send_tx_status(app_id, "armed") session.task = self._loop.run_in_executor(None, self._tx_executor_body, session) # Spawn a small watchdog that transitions armed → transmitting when # the first buffer has been consumed, and surfaces underrun / max- # duration terminations back to the hub. asyncio.create_task(self._tx_watchdog(session)) async def _handle_tx_stop(self, msg: dict) -> None: session = self._tx if session is None: return app_id = session.app_id session.stop_event.set() try: session.sdr.pause_tx() except Exception: logger.debug("pause_tx raised during stop", exc_info=True) # Wake the executor thread if it's blocked on ``queue.get``. self._drain_tx_queue(session) if session.task is not None: try: await asyncio.wait_for(asyncio.wrap_future(session.task), timeout=1.5) except asyncio.TimeoutError: logger.warning("TX executor did not exit within 1.5s after stop") except Exception: logger.debug("TX executor raised on shutdown", exc_info=True) self._close_session_sdr(session) self._tx = None await self._send_tx_status(app_id, "done") async def _handle_tx_configure(self, msg: dict) -> None: if self._tx is None: return self._tx.pending_config.update(msg.get("radio_config") or {}) # ------------------------------------------------------------------ # TX executor & watchdog def _tx_executor_body(self, session: TxSession) -> None: try: session.sdr._stream_tx(lambda n: self._tx_callback(session, n)) except Exception: logger.exception("TX stream crashed") self._schedule(self._send_tx_status(session.app_id, "error", "tx stream crashed")) def _tx_callback(self, session: TxSession, num_samples) -> np.ndarray: n = int(num_samples) # Honor stop requests: return silence one last time and let the driver # exit its loop on the next iteration (pause_tx flips _enable_tx). if session.stop_event.is_set(): return _silence(n) # Max-duration watchdog. if ( session.max_duration_s is not None and (time.monotonic() - session.started_at) >= float(session.max_duration_s) ): session.stop_event.set() try: session.sdr.pause_tx() except Exception: pass self._schedule(self._send_tx_status(session.app_id, "done", "max duration reached")) return _silence(n) # Apply queued configure at buffer boundary. if session.pending_config: cfg = session.pending_config session.pending_config = {} try: _apply_sdr_config(session.sdr, cfg) except Exception as exc: logger.debug("tx_configure apply failed: %s", exc) try: raw = session.in_queue.get(timeout=0.1) except queue.Empty: return self._underrun_fill(session, n) arr = np.frombuffer(raw, dtype=np.float32) if arr.size < 2 or arr.size % 2 != 0: logger.warning("Malformed TX frame: %d floats (must be non-zero even count)", arr.size) return self._underrun_fill(session, n) samples = (arr[0::2].astype(np.complex64) + 1j * arr[1::2].astype(np.complex64)) if samples.size < n: out = np.zeros(n, dtype=np.complex64) out[: samples.size] = samples session.last_buffer = out return out if samples.size > n: samples = samples[:n] session.last_buffer = samples if session.state == "armed": session.state = "transmitting" self._schedule(self._send_tx_status(session.app_id, "transmitting")) return samples def _underrun_fill(self, session: TxSession, n: int) -> np.ndarray: policy = session.underrun_policy if policy == "zero": return _silence(n) if policy == "repeat" and session.last_buffer is not None: buf = session.last_buffer if buf.size == n: return buf if buf.size > n: return buf[:n].copy() out = np.zeros(n, dtype=np.complex64) out[: buf.size] = buf return out # "pause" policy (default) or "repeat" before any buffer arrived. if not session.underrun_flag.is_set(): session.underrun_flag.set() session.stop_event.set() try: session.sdr.pause_tx() except Exception: pass return _silence(n) async def _tx_watchdog(self, session: TxSession) -> None: # Poll the underrun flag so we can emit status + tear down cleanly # when the callback flips the flag from the executor thread. Check # underrun_flag before stop_event, since the "pause" path sets both. while session is self._tx: if session.underrun_flag.is_set(): await self._send_tx_status(session.app_id, "underrun") await self._teardown_tx_after_underrun(session) return if session.stop_event.is_set(): return await asyncio.sleep(0.05) async def _teardown_tx_after_underrun(self, session: TxSession) -> None: if self._tx is not session: return self._drain_tx_queue(session) if session.task is not None: try: await asyncio.wait_for(asyncio.wrap_future(session.task), timeout=1.0) except asyncio.TimeoutError: logger.warning("TX executor did not exit within 1s after underrun") except Exception: logger.debug("TX executor raised during underrun teardown", exc_info=True) self._close_session_sdr(session) if self._tx is session: self._tx = None def _drain_tx_queue(self, session: TxSession) -> None: try: while True: session.in_queue.get_nowait() except queue.Empty: pass def _schedule(self, coro) -> None: loop = self._loop if loop is None: return try: asyncio.run_coroutine_threadsafe(coro, loop) except Exception: logger.debug("_schedule failed", exc_info=True) # ================================================================== # Helpers def _close_session_sdr(self, session) -> None: if session.sdr is None: return should_close = self._registry.release(session.device_key) if should_close: try: session.sdr.close() except Exception: logger.debug("SDR close raised", exc_info=True) async def _send_status(self, status: str, app_id: str) -> None: try: await self.ws.send_json({"type": "status", "status": status, "app_id": app_id}) except Exception as exc: logger.debug("Status send failed: %s", exc) async def _send_error(self, app_id: str, message: str) -> None: try: await self.ws.send_json({"type": "error", "app_id": app_id, "message": message}) except Exception as exc: logger.debug("Error-frame send failed: %s", exc) async def _send_tx_status(self, app_id: str, state: str, message: str | None = None) -> None: payload: dict = {"type": "tx_status", "app_id": app_id, "state": state} if message is not None: payload["message"] = message try: await self.ws.send_json(payload) except Exception as exc: logger.debug("tx_status 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"), "tx_sample_rate": ("tx_sample_rate",), "tx_center_frequency": ("tx_center_frequency", "tx_lo"), "tx_gain": ("tx_gain",), "tx_bandwidth": ("tx_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 _silence(num_samples: int) -> np.ndarray: """Return a ``num_samples``-length zero-filled complex64 buffer.""" return np.zeros(int(num_samples), dtype=np.complex64) 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, *, cfg: AgentConfig | None = None) -> None: """Connect to *ws_url* and run the streamer loop until cancelled.""" ws = WsClient(ws_url, token) streamer = Streamer(ws, cfg=cfg) await ws.run( streamer.on_message, streamer.build_heartbeat, on_binary=streamer.on_binary, )