Files
ha-nvidia-shield/custom_components/nvidia_shield/connection.py
T

641 lines
25 KiB
Python

"""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,
):
self.hass = hass
self.entry = entry
self.store = store if store is not None else StateStore(ttl=STATE_STORE_TTL)
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]] = []
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."""
sessions = self.active_sessions
if not sessions:
return None
if foreground_package:
for session in sessions:
if session.package == foreground_package:
return session
return sessions[0]
# ── 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()
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))
# 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)
self._cache_last_known(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()