fix: Restore the engine user's saved theme after dashboard screenshots (#1912)

* fix: Restore the engine user's saved theme after dashboard screenshots

Stock Puppet dispatches a settheme event on every cold-browser render,
and its dark query flag is presence-based, so a default screenshot
reaches the frontend as an explicit light selection. Home Assistant
persists that per-user (frontend/set_user_data, key "theme") and syncs
it to every real session of the engine token's user — flipping a
dark-mode user's web and mobile UI to light on each capture (#1909).

Bracket every capture batch with a ThemeGuard: read the engine user's
saved theme before rendering and write it back afterwards when the
render changed it. In add-on mode the guard authenticates with the
Puppet add-on's own access_token/home_assistant_url options (already
fetched during Supervisor discovery, kept in process memory only); in
sidecar/standalone/OAuth mode it falls back to ha-mcp's direct HA
credentials, which protect the user whenever both tokens belong to the
same account. Guard failures never break a capture; an attempted but
failed snapshot/restore surfaces as a tool-response warning.

Fixes #1909

* fix: Harden the theme guard per review findings

Review-toolkit and bot findings, batched:

- Surface guard warnings on the error path too: a failing batch may
  already have rendered (and clobbered the theme), so the raised
  ToolError payload now carries them, and the set-path degraded-warning
  handler merges them into the response warnings (codex P2 + toolkit).
- Gate the client-credential fallback on is_running_in_addon() instead
  of a raw SUPERVISOR_TOKEN check so embedded mode (HA core container
  env, plain admin client token) keeps theme protection (codex P2).
- Restore a never-configured baseline as {} rather than null: live
  frontend sessions ignore a null subscription push and would stay
  flipped until reload; an empty settings object re-applies default
  behavior immediately and means the same thing on the next boot
  (codex P2, verified against frontend themes-mixin).
- Carry a narrow frozen EngineCredential through EngineTarget instead
  of the raw secret-bearing Supervisor options dict; privatize the
  guard's snapshot lifecycle state; type the WebSocket handle.
- Drop the production-dead resolve_engine_url wrapper.
- Defensively tolerate malformed get_user_data responses (Gemini).
- Note the settheme-before-response ordering assumption; disambiguate
  'standalone' wording in docs/beta.md.
- Tests: guard-warning propagation into both tool responses, the
  client-credential fallback through the full capture bracket, embedded
  mode, error-path warning delivery, and malformed-response tolerance.

* fix: Close theme-guard gaps from the second review round

- Discover the Puppet add-on credential best-effort even when
  HAMCP_DASHBOARD_SCREENSHOT_ENGINE_URL is set on HA OS, so an explicit
  engine URL no longer disables the theme guard in add-on mode.
- Wait RESTORE_SETTLE_SECONDS before the post-capture read so a
  low-wait_ms render cannot race the frontend's async user-data write
  (a stale read would skip the restore and let the late write survive).
- Keep theme-guard warnings visible when image packaging fails after a
  successful render, on both the standalone and config tool paths.

---------

Co-authored-by: kingpanther13 <kingpanther13@users.noreply.github.com>

Author: kingpanther13

docs/beta.md

@@ -175,11 +175,19 @@ page and (by its design) restarts; ha-mcp surfaces this as a clear "set the
 engine's access token" error rather than a silent failure.
 
 Puppet's theme and dark-mode renderer controls dispatch Home Assistant's
-`settheme` event and can persist preferences on the frontend profile used by
-that token; even a fresh Puppet browser's first default render may save the
-default theme/dark selection. Use a dedicated Puppet account so screenshot QA
-does not alter a person's normal frontend preferences. Language selection is
-local to Puppet's browser session.
+`settheme` event, which Home Assistant persists on the frontend profile of
+the user whose token the engine runs with — and syncs to that user's real
+web and mobile sessions. Even a fresh Puppet browser's first default render
+saves a "light" selection, which used to flip a dark-mode user's whole UI to
+light on every screenshot. ha-mcp now brackets every capture: it reads that
+user's saved theme before rendering and writes it back afterwards if the
+render changed it (in add-on mode it authenticates with the Puppet add-on's
+own configured token; in sidecar / self-hosted non-add-on mode with ha-mcp's
+own HA credentials, which protects the user whenever both tokens belong to the
+same account). The restore is best-effort — a failed restore surfaces as a
+`warnings` entry on the tool response. A dedicated Puppet account remains a
+sound belt-and-suspenders setup. Language selection is local to Puppet's
+browser session.
 
 To change the Puppet engine add-on's own options (such as `keep_browser_open`)
 or to restart it, use `ha_manage_addon`; the screenshot tools only render and
@@ -254,8 +262,9 @@ render failure to a warning so it never breaks a write that already committed.
 screenshot *is* the requested payload, so a total render failure surfaces as
 an error (matching the standalone `ha_get_dashboard_screenshot` tool) rather
 than a warning a caller might miss. Because Puppet can persist theme/dark
-preferences, screenshot operations are blocked in server Read Only Mode;
-ordinary dashboard get/list/search calls remain available.
+preferences (and the theme-restore bracket writes frontend user data to undo
+that), screenshot operations are blocked in server Read Only Mode; ordinary
+dashboard get/list/search calls remain available.
 
 **Raw rendered paths remain constrained.** `ha_get_dashboard_screenshot`
 validates legacy `dashboard_path` values (rejects URLs, query strings,

src/ha_mcp/config.py

@@ -836,7 +836,7 @@ class AdvancedField(NamedTuple):
     AdvancedField("enable_websocket", "ENABLE_WEBSOCKET", bool, "operations", True),
     # Dashboard-screenshot engine URL (#1538): docker/.env users could set
     # HAMCP_DASHBOARD_SCREENSHOT_ENGINE_URL, but add-on users had no path to
-    # it. It is resolved live per capture (resolve_engine_url), so unlike the
+    # it. It is resolved live per capture (resolve_engine), so unlike the
     # time budgets it takes effect without a restart. Blank = auto-discover
     # the Puppet add-on via the Supervisor.
     AdvancedField(

src/ha_mcp/dashboard_screenshot/__init__.py

@@ -19,14 +19,18 @@
     DashboardImageCapture,
     capture_dashboard_images,
 )
-from .provision import resolve_engine_url
+from .provision import EngineTarget, resolve_engine
+from .theme_guard import EngineCredential, ThemeGuard
 
 __all__ = [
     "DEFAULT_HEIGHT",
     "DEFAULT_RENDER_TIMEOUT_SECONDS",
     "DEFAULT_WAIT_MS",
     "DEFAULT_WIDTH",
     "DashboardImageCapture",
+    "EngineCredential",
+    "EngineTarget",
+    "ThemeGuard",
     "capture_dashboard_images",
-    "resolve_engine_url",
+    "resolve_engine",
 ]

src/ha_mcp/dashboard_screenshot/capture.py

@@ -2,8 +2,11 @@
 
 The engine (balloob's Puppet add-on, or a docker-compose sidecar)
 authenticates to Home Assistant with its OWN configured long-lived token, so
-this client passes only the dashboard path + render parameters — no HA token
-ever flows through ha-mcp or the LLM for screenshots.
+this client passes only the dashboard path + render parameters to the engine
+— no HA token is ever sent to the engine or exposed to the LLM. The
+:mod:`theme_guard` bracket around each capture batch may hold the engine's
+token in server memory to restore the engine user's saved theme (issue
+#1909); see that module's docstring for the containment rules.
 """
 
 from __future__ import annotations
@@ -20,7 +23,8 @@
 
 from ..errors import ErrorCode, create_error_response
 from ..tools.helpers import raise_tool_error
-from .provision import TOKEN_HINT, resolve_engine_url
+from .provision import TOKEN_HINT, resolve_engine
+from .theme_guard import ThemeGuard
 
 logger = logging.getLogger(__name__)
 
@@ -757,9 +761,30 @@ async def _request_or_collect_failure(
         return None
 
 
+def _viewport_params(
+    options: _CaptureOptions, viewport: _ViewportRequest
+) -> dict[str, str]:
+    """Build one viewport's engine query parameters."""
+    params: dict[str, str] = {
+        "viewport": f"{viewport.width}x{viewport.height}",
+        "zoom": str(options.zoom),
+        "wait": str(options.wait_ms),
+        "format": options.image_format,
+    }
+    if options.theme is not None:
+        params["theme"] = options.theme
+    if options.dark_mode:
+        # Puppet checks only for query-key presence.
+        params["dark"] = ""
+    if options.language is not None:
+        params["lang"] = options.language
+    return params
+
+
 def _complete_capture_batch(
     captures: list[DashboardImageCapture],
     partial_failures: list[dict[str, Any]] | None,
+    guard_warnings: list[str],
 ) -> list[DashboardImageCapture]:
     """Return partial successes, but raise when every requested capture failed."""
     if not captures and partial_failures:
@@ -769,10 +794,41 @@ def _complete_capture_batch(
             "failure_count": len(partial_failures),
             "screenshot_failures": partial_failures,
         }
+        if guard_warnings:
+            aggregate_failure["warnings"] = [
+                *aggregate_failure.get("warnings", []),
+                *guard_warnings,
+            ]
         raise_tool_error(aggregate_failure)
     return captures
 
 
+def _reraise_with_guard_warnings(error: ToolError, warnings: list[str]) -> NoReturn:
+    """Re-raise a capture ToolError with the theme guard's warnings attached.
+
+    A batch that raises may already have rendered (and clobbered the engine
+    user's theme) before failing, so a failed restore must stay visible on
+    the error path too — otherwise the tool contract ("an attempted but
+    failed restore surfaces as a warning") silently breaks exactly when the
+    theme is most likely left flipped.
+    """
+    if not warnings:
+        raise error
+    try:
+        payload = json.loads(str(error))
+    except (json.JSONDecodeError, TypeError):
+        payload = None
+    if not isinstance(payload, dict):
+        raise error
+    existing = payload.get("warnings")
+    payload["warnings"] = [
+        *(existing if isinstance(existing, list) else []),
+        *warnings,
+    ]
+    raise_tool_error(payload)
+    raise AssertionError("unreachable: raise_tool_error always raises")
+
+
 async def capture_dashboard_images(
     dashboard_path: str,
     *,
@@ -789,13 +845,21 @@ async def capture_dashboard_images(
     image_format: ScreenshotFormat = "png",
     render_timeout_seconds: float = DEFAULT_RENDER_TIMEOUT_SECONDS,
     partial_failures: list[dict[str, Any]] | None = None,
+    client: Any | None = None,
+    capture_warnings: list[str] | None = None,
 ) -> list[DashboardImageCapture]:
     """Render one or more ordered dashboard images via the screenshot engine.
 
     Named presets override explicit dimensions. An explicit orientation swaps
     preset or custom dimensions when needed; omitting it preserves each
     preset's native orientation. ``full_page`` is a compatibility alias for
     requesting the engine's native ``WIDTHxauto`` viewport.
+
+    The batch is bracketed by a :class:`ThemeGuard` that restores the engine
+    user's saved frontend theme if rendering changed it (issue #1909).
+    ``client`` supplies the guard's non-add-on credential fallback and
+    ``capture_warnings`` collects the guard's non-fatal warnings; both are
+    optional and never affect the captures themselves.
     """
     path = _validate_dashboard_path(dashboard_path)
     options = validate_capture_parameters(
@@ -814,94 +878,102 @@ async def capture_dashboard_images(
     )
     viewports = _capture_viewports(options)
 
-    engine = await resolve_engine_url()
+    engine_target = await resolve_engine()
+    engine = engine_target.url
     url = f"{engine}/{path}"
     mime_type = _MIME_TYPES[options.image_format]
     captures: list[DashboardImageCapture] = []
 
-    async with httpx.AsyncClient(
-        timeout=httpx.Timeout(options.render_timeout_seconds)
-    ) as http_client:
-        for capture_index, viewport in enumerate(viewports):
-            params: dict[str, str] = {
-                "viewport": f"{viewport.width}x{viewport.height}",
-                "zoom": str(options.zoom),
-                "wait": str(options.wait_ms),
-                "format": options.image_format,
-            }
-            if options.theme is not None:
-                params["theme"] = options.theme
-            if options.dark_mode:
-                # Puppet checks only for query-key presence.
-                params["dark"] = ""
-            if options.language is not None:
-                params["lang"] = options.language
-
-            request_context = {
-                "path": path,
-                "preset": viewport.preset,
-                "width": viewport.width,
-                "height": viewport.height,
-                "requested_format": options.image_format,
-                "capture_index": capture_index,
-                "capture_count": len(viewports),
-                "completed_count": len(captures),
-            }
-            aggregate_bytes = sum(capture.size_bytes for capture in captures)
-            remaining_batch_bytes = MAX_BATCH_PAYLOAD_BYTES - aggregate_bytes
-            if remaining_batch_bytes <= 0:
-                assert partial_failures is not None
-                partial_failures.append(
-                    create_error_response(
-                        ErrorCode.IMAGE_PAYLOAD_TOO_LARGE,
-                        "Screenshot image batch reached the server's safe "
-                        "inline-image limit.",
-                        context={
-                            **request_context,
-                            "aggregate_bytes_before_capture": aggregate_bytes,
-                            "limit_kind": "batch",
-                            "limit_bytes": MAX_BATCH_PAYLOAD_BYTES,
-                        },
+    # The engine dispatches a theme write into the authenticated frontend on
+    # cold renders, which Home Assistant persists to the engine user's real
+    # profile (issue #1909). Snapshot before, restore after — including when
+    # the batch fails, since the engine may already have rendered (and
+    # written) before the failure.
+    guard = ThemeGuard.for_capture(engine_target.addon_credential, client)
+    await guard.take_snapshot()
+    batch_error: ToolError | None = None
+    try:
+        async with httpx.AsyncClient(
+            timeout=httpx.Timeout(options.render_timeout_seconds)
+        ) as http_client:
+            for capture_index, viewport in enumerate(viewports):
+                params = _viewport_params(options, viewport)
+
+                request_context = {
+                    "path": path,
+                    "preset": viewport.preset,
+                    "width": viewport.width,
+                    "height": viewport.height,
+                    "requested_format": options.image_format,
+                    "capture_index": capture_index,
+                    "capture_count": len(viewports),
+                    "completed_count": len(captures),
+                }
+                aggregate_bytes = sum(capture.size_bytes for capture in captures)
+                remaining_batch_bytes = MAX_BATCH_PAYLOAD_BYTES - aggregate_bytes
+                if remaining_batch_bytes <= 0:
+                    assert partial_failures is not None
+                    partial_failures.append(
+                        create_error_response(
+                            ErrorCode.IMAGE_PAYLOAD_TOO_LARGE,
+                            "Screenshot image batch reached the server's safe "
+                            "inline-image limit.",
+                            context={
+                                **request_context,
+                                "aggregate_bytes_before_capture": aggregate_bytes,
+                                "limit_kind": "batch",
+                                "limit_bytes": MAX_BATCH_PAYLOAD_BYTES,
+                            },
+                        )
                     )
-                )
-                break
-
-            outcome = await _request_or_collect_failure(
-                http_client,
-                url=url,
-                path=path,
-                engine=engine,
-                params=params,
-                options=options,
-                viewport=viewport,
-                mime_type=mime_type,
-                request_context=request_context,
-                aggregate_bytes=aggregate_bytes,
-                remaining_batch_bytes=remaining_batch_bytes,
-                partial_failures=partial_failures,
-            )
-            if outcome is None:
-                if (
-                    partial_failures
-                    and partial_failures[-1].get("limit_kind") == "batch"
-                ):
                     break
-                continue
-            image_data, capture_height, fallback_used = outcome
-
-            captures.append(
-                DashboardImageCapture(
-                    data=image_data,
-                    width=viewport.width,
-                    height=capture_height,
-                    preset=viewport.preset,
-                    orientation=viewport.orientation,
-                    image_format=options.image_format,
-                    mime_type=mime_type,
-                    size_bytes=len(image_data),
+
+                outcome = await _request_or_collect_failure(
+                    http_client,
+                    url=url,
+                    path=path,
+                    engine=engine,
+                    params=params,
                     options=options,
-                    legacy_full_page_fallback=fallback_used,
+                    viewport=viewport,
+                    mime_type=mime_type,
+                    request_context=request_context,
+                    aggregate_bytes=aggregate_bytes,
+                    remaining_batch_bytes=remaining_batch_bytes,
+                    partial_failures=partial_failures,
                 )
-            )
-
-    return _complete_capture_batch(captures, partial_failures)
+                if outcome is None:
+                    if (
+                        partial_failures
+                        and partial_failures[-1].get("limit_kind") == "batch"
+                    ):
+                        break
+                    continue
+                image_data, capture_height, fallback_used = outcome
+
+                captures.append(
+                    DashboardImageCapture(
+                        data=image_data,
+                        width=viewport.width,
+                        height=capture_height,
+                        preset=viewport.preset,
+                        orientation=viewport.orientation,
+                        image_format=options.image_format,
+                        mime_type=mime_type,
+                        size_bytes=len(image_data),
+                        options=options,
+                        legacy_full_page_fallback=fallback_used,
+                    )
+                )
+    except ToolError as exc:
+        # Held (not re-raised here) so the restore in ``finally`` runs first
+        # and its outcome can be attached to the error payload below.
+        batch_error = exc
+    finally:
+        await guard.restore()
+        if capture_warnings is not None:
+            capture_warnings.extend(guard.warnings)
+
+    if batch_error is not None:
+        _reraise_with_guard_warnings(batch_error, guard.warnings)
+    return _complete_capture_batch(captures, partial_failures, guard.warnings)