"""ShieldConnectionManager — entry-scoped owner of the ADB connection (Phase 1). The ADB socket, its watchdog and all ADB-derived state live here instead of on the media-player entity (plan blockers B1/B2/B5): - one persistent socket; a watchdog task reconnects with exponential backoff (5 s doubling to a 300 s cap, ±25 % jitter) when it drops, and pulls the heartbeat data (wake + media sessions + running apps) while connected; - ADB-derived values live in a :class:`StateStore` with per-field timestamps and TTL expiry — stale values read as expired (None/empty) instead of clinging on; - entities are thin callback-views: they ``register_listener()`` and write HA state when the store changes (the androidtv core ``entity.py`` pattern); - ADB is optional: when no poller factory is provided the manager is *passive* and the entry still sets up — a failed ADB connect is a background retry, never a ``ConfigEntryNotReady``. This module deliberately stays Home-Assistant-free at import time so it is unit-testable with plain pytest; every hass capability is injected (the entry builds the concrete poller factory). Task creation additionally works without hass (plain ``asyncio.create_task``) for the unit tests. """ import asyncio import contextlib import logging import random import time from collections.abc import Callable from typing import Any from .adb_session import AdbMediaSessionPoller from .const import ( ADB_POLL_INTERVAL, ADB_RECONNECT_BACKOFF_CAP, ADB_RECONNECT_BACKOFF_START, ADB_RECONNECT_JITTER, STATE_STORE_TTL, SYSTEM_MEDIA_PACKAGES, ) from .ha_adb_common import MediaSessionInfo, parse_audio_playback _LOGGER = logging.getLogger(__name__) # StateStore field names FIELD_AWAKE = "awake" FIELD_SESSIONS = "sessions" FIELD_AUDIO_PLAYBACK = "audio_playback" FIELD_RUNNING_APPS = "running_apps" # Protocol transport (androidtvremote2) push fields FIELD_IS_ON = "is_on" FIELD_IS_AVAILABLE = "is_available" FIELD_CURRENT_APP = "current_app" FIELD_VOLUME_INFO = "volume_info" # ADB fallback foreground identity (P3-2): stamped only when the protocol # push is silent; consumers prefer FIELD_CURRENT_APP over this. FIELD_FOREGROUND = "foreground" # ADB-derived fields (cleared when the ADB connection goes down — they are # regenerated by the watchdog); protocol pushes (IS_ON/IS_AVAILABLE/ # CURRENT_APP/VOLUME_INFO) are device-reported truth and expire by TTL only, # so an ADB outage never clobbers them. ADB_FIELDS = ( FIELD_AWAKE, FIELD_SESSIONS, FIELD_AUDIO_PLAYBACK, FIELD_RUNNING_APPS, FIELD_FOREGROUND, ) PROTOCOL_FIELDS = ( FIELD_IS_ON, FIELD_IS_AVAILABLE, FIELD_CURRENT_APP, FIELD_VOLUME_INFO, ) ALL_FIELDS = ADB_FIELDS + PROTOCOL_FIELDS def volume_tuple(info) -> tuple | None: """Normalize the protocol's volume push to ``(level, max, muted)``. The library delivers a VolumeInfo TypedDict (``level``/``max``/``muted``). Readout only — there is no protocol-native VOLUME_SET (C-VOL), so nothing here ever sends one. """ if info is None: return None try: return ( float(info.get("level", 0.0)), float(info.get("max", 1.0)), bool(info.get("muted", False)), ) except (AttributeError, TypeError, ValueError): return None class StateStore: """Timed per-field state cache shared by all entities of one entry. Reading an expired (or missing) field returns the caller-supplied default, so a stale "playing" session can never outlive the TTL (plan blocker B5). Values are written only via :meth:`set`, which re-stamps the field's freshness timestamp. """ def __init__(self, ttl: float = STATE_STORE_TTL, clock: Callable[[], float] = time.monotonic): self._ttl = ttl self._clock = clock self._values: dict[str, tuple[float, Any]] = {} def set(self, name: str, value: Any) -> None: """Store a value, re-stamping its freshness timestamp.""" self._values[name] = (self._clock(), value) def get(self, name: str, default: Any = None) -> Any: """Return the field's value, or ``default`` when it is missing/expired.""" stamped = self._values.get(name) if stamped is None: return default at, value = stamped if self._clock() - at > self._ttl: return default return value def clear(self, *names: str) -> None: """Drop fields immediately so they read as expired on the next get.""" for name in names: self._values.pop(name, None) @property def ttl(self) -> float: return self._ttl class ShieldConnectionManager: """Owner of one entry's persistent connections (both transports) + state. Transport A — ADB (adb_shell): the watchdog keeps a persistent socket, pulls heartbeats and reconnects with its own backoff ladder. Transport B — remote protocol (androidtvremote2, Phase 2): the library owns reconnection (``keep_reconnecting`` with internal backoff), pushes ``is_on``/``is_available``/``current_app``/``volume_info`` into the state store, and surfaces invalid certs (device hard reset) via a reauth hook. :param poller_factory: builds the ADB poller when the watchdog starts; skipped entirely for entries without ADB (passive manager) or tests that inject a ready-to-use ``poller`` directly. :param protocol_factory: builds the AndroidTVRemote client; instead pass ``protocol_settings`` (host/client_name/certfile/keyfile) and the manager builds it, or ``protocol`` directly (tests). :param invalid_auth_callback: called when the device rejects our certificates (hard reset) — the entry-owned reauth hook (t9). """ def __init__( self, hass=None, entry=None, poller_factory: Callable[[], AdbMediaSessionPoller] | None = None, poller: AdbMediaSessionPoller | None = None, store: StateStore | None = None, protocol_factory: Callable[[], Any] | None = None, protocol: Any = None, protocol_settings: dict | None = None, invalid_auth_callback: Callable[[], None] | None = None, clock: Callable[[], float] = time.monotonic, ): self.hass = hass self.entry = entry self.store = ( store if store is not None else StateStore(ttl=STATE_STORE_TTL, clock=clock) ) self._poller_factory = poller_factory self._poller = poller # Remote-protocol transport (androidtvremote2): entry passes either a # settings dict (manager itself builds the real client) or a factory / # ready-made client (tests, or the config flow when it pre-pairs). self._protocol_factory = protocol_factory self._protocol = protocol self._protocol_settings = protocol_settings self._invalid_auth_callback = invalid_auth_callback self._protocol_loop: asyncio.AbstractEventLoop | None = None self._listeners: list[Callable[[], None]] = [] # P3-2d last-known-state cache: package → (session info, stamp). # Serves the foreground app's last known playback state when it # publishes no MediaSession (TiviMate, R2). TTL-capped: 'playing' # is NEVER resurrected past STATE_STORE_TTL. self._last_known: dict[str, tuple[MediaSessionInfo, float]] = {} # Injectable clock keeps the TTL expiry tests deterministic. self._clock = clock self._tasks: list[asyncio.Task] = [] self._stopping = False self._failures = 0 self._down_since = False # listeners already told the connection is down # ── One-shot connection test (pairing UX) ── @staticmethod async def test_connect(host: str, port: int, key_path: str) -> bool: """Single-shot ADB connect used by the config-flow pairing step. Throws away the poller immediately — the config flow runs before any entry exists, so it must never touch an entry's persistent socket. """ poller = AdbMediaSessionPoller(host, port, key_path) try: return await poller.connect() finally: await poller.disconnect() # ── Remote-protocol transport (androidtvremote2, Phase 2) ── @staticmethod def build_protocol_client( host: str, client_name: str, certfile: str, keyfile: str, api_port: int = 6466, pair_port: int = 6467, loop=None, ): """Build a real androidtvremote2 client (cert/PIN flow hook, t9). Kept as a manager-level factory so the config flow never hand-rolls the constructor; the import stays local so this module remains importable where androidtvremote2 is not installed. """ try: import androidtvremote2 except ImportError: _LOGGER.warning( "androidtvremote2 is not installed — the remote-protocol " "transport is disabled and the entry degrades to ADB-only." ) return None return androidtvremote2.AndroidTVRemote( client_name=client_name, certfile=certfile, keyfile=keyfile, host=host, api_port=api_port, pair_port=pair_port, loop=loop, ) def _build_protocol(self): """Default factory from the entry settings (host + cert/key paths).""" settings = self._protocol_settings or {} return self.build_protocol_client( host=settings["host"], client_name=settings.get("client_name", "Home Assistant"), certfile=settings["certfile"], keyfile=settings["keyfile"], ) async def _start_protocol(self) -> None: """Spin up the protocol transport alongside the ADB watchdog. One owner per transport: the library's own ``keep_reconnecting()`` handles reconnection with its internal backoff, so no competing watchdog is built here. Push callbacks write straight into the StateStore and ping registered listeners. """ if self._protocol is None: if self._protocol_factory is None: if self._protocol_settings is None: return # protocol transport disabled for this entry self._protocol_factory = self._build_protocol self._protocol = self._protocol_factory() client = self._protocol if client is None: return # Real client: make sure local certificate material exists before it # connects (fails soft — pair flows can regenerate). generate_cert = getattr(client, "async_generate_cert_if_missing", None) if generate_cert is not None: with contextlib.suppress(Exception): await generate_cert() self._protocol_loop = asyncio.get_running_loop() client.add_is_on_updated_callback(self._on_is_on) client.add_is_available_updated_callback(self._on_is_available) client.add_current_app_updated_callback(self._on_current_app) client.add_volume_info_updated_callback(self._on_volume_info) # One owner per transport: the library's keep_reconnecting() creates # and owns its reconnect task internally — we never wrap or duplicate # it (the stop path closes the client, which tears the task down). client.keep_reconnecting(self._on_invalid_auth) def _on_is_on(self, on) -> None: self._protocol_push(FIELD_IS_ON, bool(on)) def _on_is_available(self, is_available) -> None: self._protocol_push(FIELD_IS_AVAILABLE, bool(is_available)) def _on_current_app(self, app_package) -> None: self._protocol_push(FIELD_CURRENT_APP, app_package) def _on_volume_info(self, info) -> None: self._protocol_push(FIELD_VOLUME_INFO, volume_tuple(info)) def _on_invalid_auth(self) -> None: """Device hard reset — certificates rejected. The library keeps its own backoff, but retrying can never succeed since certificated auth is gone — the entry-level callback (reauth flow) decides the user-facing response (t9). """ _LOGGER.warning( "Shield remote-protocol auth is invalid (device hard reset?) — " "notifying the entry's reauth hook" ) if self._invalid_auth_callback is not None: try: self._invalid_auth_callback() except Exception: _LOGGER.exception("Shield invalid-auth callback failed") def _protocol_push(self, field: str, value) -> None: """Stamp a push value into the store and ping listeners. Push callbacks may arrive off the event loop (grpc worker threads); route the listener notification back onto the captured loop. """ self.store.set(field, value) loop = self._protocol_loop if loop is not None and loop.is_running(): try: loop.call_soon_threadsafe(self._notify) return except RuntimeError: pass # loop closed during shutdown — skip the notify self._notify() # ── Introspection (thin views for entities) ── @property def available(self) -> bool: """True when ADB is active for this entry (a poller exists).""" return self._poller is not None @property def connected(self) -> bool: return bool(self._poller and self._poller.connected) @property def awake(self) -> bool | None: """Device wake state (exact mWakefulness parse), None = unknown.""" return self.store.get(FIELD_AWAKE) @property def active_sessions(self) -> list[MediaSessionInfo]: """Stored media sessions, sorted playing > paused > other.""" return self.store.get(FIELD_SESSIONS, []) @property def audio_playback(self) -> list[str]: """Packages that recently played audio, ordered last-played-first.""" return self.store.get(FIELD_AUDIO_PLAYBACK, []) @property def running_apps(self) -> list[dict]: """Visible recent-task apps, ordered by recency.""" return self.store.get(FIELD_RUNNING_APPS, []) @property def protocol_active(self) -> bool: """True when the remote-protocol transport is wired for this entry. Entities use this instead of reaching into ``_protocol``: a passive manager means "no transport", not "broken transport". """ return self._protocol is not None @property def is_on(self) -> bool | None: """Protocol power push (device reports on/off state).""" return self.store.get(FIELD_IS_ON) @property def is_available(self) -> bool | None: """Protocol availability push (device reachable at the protocol).""" return self.store.get(FIELD_IS_AVAILABLE) @property def current_app(self) -> str | None: """Protocol current-app push (raw package; name resolution deferred).""" return self.store.get(FIELD_CURRENT_APP) @property def foreground_app(self) -> str | None: """Best foreground identity: protocol push first, ADB parse fallback (P3-2 store priority: FIELD_CURRENT_APP > FIELD_FOREGROUND).""" return self.store.get(FIELD_CURRENT_APP) or self.store.get(FIELD_FOREGROUND) @property def volume_info(self) -> tuple | None: """Protocol volume push as ``(level, max, muted)`` — readout only (no protocol VOLUME_SET exists, C-VOL).""" return self.store.get(FIELD_VOLUME_INFO) def best_session( self, foreground_package: str | None = None ) -> MediaSessionInfo | None: """Best media session: foreground app first, then playing > paused. P3-2d: when there is no live session at all, the foreground app's last known state is served from the TTL-capped cache (TiviMate publishes no MediaSession — RESEARCH_SHIELD R2). """ sessions = self.active_sessions if foreground_package: for session in sessions: if session.package == foreground_package: return session if sessions: return sessions[0] if foreground_package: return self.last_known_session(foreground_package) return None def _cache_last_known_sessions(self) -> None: """Stamp every live session's state as the package's last known.""" now = self._clock() for session in self.active_sessions: if session.package: self._last_known[session.package] = (session, now) def last_known_session(self, package: str) -> MediaSessionInfo | None: """Last known session state for a package, TTL-capped. Returns None once STATE_STORE_TTL passes — 'playing' must never be resurrected from a stale cache (P3-2d expiry contract). """ cached = self._last_known.get(package) if cached is None: return None info, stamp = cached if self._clock() - stamp > STATE_STORE_TTL: self._last_known.pop(package, None) return None return info def last_known_map(self) -> dict[str, MediaSessionInfo]: """TTL-filtered snapshot of the per-package last-known cache.""" result: dict[str, MediaSessionInfo] = {} for package in tuple(self._last_known): info = self.last_known_session(package) if info is not None: result[package] = info return result def last_known_state( self, package: str ) -> tuple[MediaSessionInfo | None, float | None]: """(last known session, freshness stamp) for a package — the stamp lets callers age the served state in their UI.""" cached = self._last_known.get(package) if cached is None: return None, None info, stamp = cached if self._clock() - stamp > STATE_STORE_TTL: self._last_known.pop(package, None) return None, None return info, stamp # ── Listener API (entities are thin callback-views) ── def register_listener(self, listener: Callable[[], None]) -> Callable[[], None]: """Register a refresh callback; returns an unsubscribe callable.""" self._listeners.append(listener) def _unregister() -> None: with contextlib.suppress(ValueError): self._listeners.remove(listener) return _unregister def _notify(self) -> None: for listener in tuple(self._listeners): try: listener() except Exception: _LOGGER.exception("Shield state listener failed") # ── Delegation (commands over the managed socket) ── async def send_media_command(self, command: str) -> bool: """Send a media key command over the persistent socket.""" if not self.available: return False return await self._poller.send_media_command(command) async def bring_to_foreground(self, package: str) -> bool: if not self.available: return False return await self._poller.bring_to_foreground(package) # -- Protocol command delegation (Phase 2) -- async def send_protocol_key(self, key: str, direction: str = "SHORT") -> bool: """Send a remote-protocol key push (fast, non-blocking). Returns False when the transport is down or the send raises — entities then fall back to the ADB keyevent path. """ if self._protocol is None: return False try: self._protocol.send_key_command(key, direction) return True except Exception as err: _LOGGER.debug("Protocol key %s failed: %s", key, err) return False async def launch_app(self, app_link_or_app_id: str) -> bool: """Launch an app or URL over the protocol transport. ``send_launch_app_command`` auto-prefixes ``market://launch?id=`` for bare package ids (C-LAUNCH); deep links pass through unchanged. """ if self._protocol is None: return False try: self._protocol.send_launch_app_command(app_link_or_app_id) return True except Exception as err: _LOGGER.debug("Protocol launch %s failed: %s", app_link_or_app_id, err) return False async def send_digit_key(self, digit: str) -> bool: """Channel digit via the ADB keyevent path (fallback for §7).""" if not self.available: return False return await self._poller.send_digit_keyevent(digit) async def send_enter_key(self) -> bool: """ENTER via the ADB keyevent path (KEYCODE_ENTER = 66, §7 fallback).""" if not self.available: return False return await self._poller.send_enter_keyevent() # ── Lifecycle ── async def async_start(self) -> None: """Start the persistent connection loop. Never validates the ADB connection here: setup must not fail because of ADB (it is optional), so all failure handling lives in the watchdog's background retry. """ self._stopping = False self._tasks.append(self._create_task(self._watchdog(), "shield_adb_watchdog")) await self._start_protocol() async def async_stop(self) -> None: """Cancel all manager tasks and close the socket. Called on entry unload and on EVENT_HOMEASSISTANT_STOP. """ self._stopping = True tasks, self._tasks = self._tasks, [] for task in tasks: task.cancel() self._listeners.clear() self._last_known.clear() if self._poller is not None: with contextlib.suppress(Exception): await self._poller.disconnect() self._poller = None if self._protocol is not None: with contextlib.suppress(Exception): self._protocol.disconnect() self._protocol = None def _create_task(self, coro, name: str) -> asyncio.Task: if self.hass is not None: return self.hass.async_create_background_task(coro, name) return asyncio.create_task(coro, name=name) async def _watchdog(self) -> None: """Persistent loop: reconnect with backoff when down, refresh when up. A passive manager (no poller possible) exits immediately — the entry recreates it when options change (full reload). """ try: if self._poller is None: if self._poller_factory is None: return self._poller = self._poller_factory() if self._poller is None: return backoff = ADB_RECONNECT_BACKOFF_START while not self._stopping: delay = await self._watchdog_cycle(backoff) if self._stopping: break if self.connected: backoff = ADB_RECONNECT_BACKOFF_START else: backoff = min(backoff * 2, ADB_RECONNECT_BACKOFF_CAP) await asyncio.sleep(delay) except asyncio.CancelledError: raise except Exception: _LOGGER.exception("Shield connection watchdog crashed") async def _watchdog_cycle(self, backoff: float) -> float: """Run one watchdog cycle; return seconds to sleep before the next. When the connection is up this is the heartbeat: wake check + (when awake) media sessions + running apps, all stored with timestamps. The current backoff sleeps after failures; success resets it.""" try: if not self.connected and not await self._poller.connect(): await self._note_failure(backoff) return self._jitter(backoff) if self._failures: self._failures = 0 self._down_since = False _LOGGER.info("ADB connection recovered") await self._refresh() return ADB_POLL_INTERVAL except Exception as err: await self._note_failure(backoff) _LOGGER.debug("Watchdog cycle failed: %s", err) return self._jitter(backoff) async def _note_failure(self, backoff: float) -> None: """Bookkeep a failed cycle: log sparingly, drop stale store, notify.""" self._failures += 1 if self._failures == 1 or self._failures % 6 == 0: _LOGGER.warning( "ADB connection down (%s consecutive failures) — retrying in ~%.0fs", self._failures, backoff, ) if not self._down_since: self._down_since = True # Reconnect decision (health_check's documented role): close the # socket so the next cycle rebuilds it from scratch instead of # trusting shell()'s auto-reconnect on a broken transport. try: await self._poller.disconnect() except Exception: _LOGGER.debug("disconnect during reconnect decision failed", exc_info=True) # Only ADB-derived fields expire here: protocol pushes are # device-reported truth and must survive an ADB outage (t11 F2). self.store.clear(*ADB_FIELDS) self._notify() @staticmethod def _jitter(delay: float) -> float: """Backoff with ±jitter so many clients do not reconnect in lockstep.""" return max( 1.0, delay * (1.0 + random.uniform(-ADB_RECONNECT_JITTER, ADB_RECONNECT_JITTER)), ) async def _refresh(self) -> None: """Heartbeat: pull all ADB-derived state into the store (with stamps).""" poller = self._poller awake = await poller.is_awake() if awake is None: # Failed poll (connection dropped / device unreachable) — surface # as a cycle failure instead of clinging to the last values. raise ConnectionError("ADB wake poll failed") self.store.set(FIELD_AWAKE, awake) if awake: # Full dump, no truncation (B4) — low sessions must not be lost. await poller.poll_media_sessions(head_lines=None) self.store.set(FIELD_SESSIONS, list(poller.all_sessions)) self._cache_last_known_sessions() # Zero-cost hint from the same dump; system processes filtered # (mediashell/systemui/tv settings/diagtools report "audio"). playback = [ pkg for pkg in parse_audio_playback(poller.last_dump or "") if pkg not in SYSTEM_MEDIA_PACKAGES ] self.store.set(FIELD_AUDIO_PLAYBACK, playback) self.store.set(FIELD_RUNNING_APPS, await poller.poll_running_apps()) # Foreground fallback (P3-2): ADB parse only when the protocol # push is silent — store priority is protocol > ADB, so this # never overwrites a live push. if self.store.get(FIELD_CURRENT_APP) is None: foreground = await poller.poll_foreground() if foreground: self.store.set(FIELD_FOREGROUND, foreground) else: # Standby: playback data is not current — expire it. self.store.set(FIELD_SESSIONS, []) self.store.set(FIELD_AUDIO_PLAYBACK, []) self.store.set(FIELD_RUNNING_APPS, []) self._notify()