fix(screenshot): re-enable the theme guard around dashboard captures

The ThemeGuard bracket in capture_dashboard_images was commented out on
the premise that upstream Puppet had fixed the cold-render settheme
dispatch that made it necessary (#1909). That fix never shipped:
balloob/home-assistant-addons#89 is merged to that repo's master but
unreleased -- it landed after the 2.6.0 version bump (#88), and
puppet/config.yaml still reads 2.6.0 with no settheme entry in its
CHANGELOG. Every deployment on a released Puppet therefore still has the
theme persisted onto the engine token user's profile and synced to that
user's live web and mobile sessions, with nothing restoring it.

The upstream fix is also narrower than the disable assumed. Its title is
exact -- "don't dispatch settheme when no theme/dark was requested" -- so
a capture passing theme= or dark_mode=true still writes even once it
ships. The bracket is therefore restored unconditionally rather than
gated on engine version.

ha_get_dashboard_screenshot already documents the restore ("The engine
user's saved theme preference is restored after the capture (best
effort)"), so the disabled bracket also left the tool description
promising behaviour the code did not perform.

TestCaptureBracketDisabled asserted the bracket stayed off; it is
inverted to TestCaptureBracket, which now asserts the snapshot/restore
sessions open and the clobbered theme is written back, including on the
capture-failure path.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_012Nm7tyA1nfxNCWFXaR3AxV

Author: kingpanther13

docs/beta.md

@@ -177,18 +177,18 @@ token grants. If the token is missing or invalid, Puppet lands on the login
 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 used to dispatch Home
-Assistant's `settheme` event on every cold render, which Home Assistant
-persisted on the frontend profile of the user whose token the engine runs with
-— and synced to that user's real web and mobile sessions, flipping a dark-mode
-user's whole UI to light on every screenshot (#1909). Recent Puppet versions
-fixed that cold-render dispatch, so ha-mcp's snapshot/restore bracket around
-each capture is now disabled (#1991); the guard code is retained so it can be
-switched back on if a future engine regression reintroduces the write. If you
-run an older Puppet build, update the app (or your self-hosted sidecar
-image) — older engines still persist the theme selection and will keep
-flipping it. A dedicated Puppet account remains a sound belt-and-suspenders
-setup. Language selection is local to Puppet's browser session.
+Puppet's theme and dark-mode renderer controls dispatch Home Assistant's
+`settheme` event on render, 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, flipping that user's whole UI on every
+screenshot (#1909). ha-mcp brackets each capture batch with a snapshot/restore
+of that user's saved theme, so the flip is undone automatically. The upstream
+Puppet fix for the dispatch (balloob/home-assistant-addons#89) is merged but
+unreleased as of Puppet 2.6.0, and is scoped to renders that request no
+`theme`/`dark` — an explicit theme or dark-mode capture still writes even once
+it ships, so the bracket stays on regardless of engine version. 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 app's own options (such as `keep_browser_open`)
 or to restart it, use `ha_manage_app`; the screenshot tools only render and

src/ha_mcp/dashboard_screenshot/capture.py

@@ -855,15 +855,9 @@ async def capture_dashboard_images(
     preset's native orientation. ``full_page`` is a compatibility alias for
     requesting the engine's native ``WIDTHxauto`` viewport.
 
-    The batch used to be bracketed by a :class:`ThemeGuard` that restored the
-    engine user's saved frontend theme when a cold render changed it (issue
-    #1909). That bracket is currently disabled (#1991): upstream Puppet no
-    longer dispatches ``settheme`` on cold renders, so there is nothing to undo.
-    The guard construction and the ``client`` / ``capture_warnings`` plumbing
-    are retained so the bracket can be re-enabled by uncommenting the
-    snapshot/restore calls if a future engine regression reintroduces the write;
-    while disabled ``capture_warnings`` simply stays empty and never affects the
-    captures themselves.
+    The batch is bracketed by a :class:`ThemeGuard` that restores the engine
+    user's saved frontend theme when a render changes it (issue #1909). Guard
+    failures are non-fatal and surface through ``capture_warnings``.
     """
     path = _validate_dashboard_path(dashboard_path)
     options = validate_capture_parameters(
@@ -888,16 +882,20 @@ async def capture_dashboard_images(
     mime_type = _MIME_TYPES[options.image_format]
     captures: list[DashboardImageCapture] = []
 
-    # ThemeGuard bracket — currently DISABLED (#1991). Stock Puppet used to
-    # dispatch a theme write into the authenticated frontend on cold renders,
-    # which Home Assistant persisted to the engine user's real profile (#1909);
-    # ha-mcp snapshotted before and restored after to undo it. Upstream Puppet
-    # has since fixed the cold-render settheme dispatch, so the bracket is no
-    # longer needed. The guard is still constructed (and the snapshot/restore
-    # calls kept below, commented out) so it can be re-enabled by uncommenting
-    # if a future engine regression reintroduces the write.
+    # ThemeGuard bracket. Puppet dispatches a ``settheme`` event into the
+    # authenticated frontend, which Home Assistant persists to the engine
+    # user's real profile and syncs to that user's live sessions (#1909).
+    # ha-mcp snapshots the saved theme before the batch and restores it after.
+    #
+    # The bracket was previously disabled on the premise that upstream Puppet
+    # had fixed the dispatch. It had not shipped: the fix
+    # (balloob/home-assistant-addons#89) is merged to that repo's master but
+    # unreleased — the newest Puppet release, 2.6.0, predates it. It is also
+    # scoped to the no-parameter case ("don't dispatch settheme when no
+    # theme/dark was requested"), so an explicit ``theme=``/``dark`` capture
+    # still writes even once it does ship.
     guard = ThemeGuard.for_capture(engine_target.addon_credential, client)
-    # await guard.take_snapshot()
+    await guard.take_snapshot()
     batch_error: ToolError | None = None
     try:
         async with httpx.AsyncClient(
@@ -977,7 +975,7 @@ async def capture_dashboard_images(
         # and its outcome can be attached to the error payload below.
         batch_error = exc
     finally:
-        # await guard.restore()  # ThemeGuard bracket disabled (#1991) — see above.
+        await guard.restore()
         if capture_warnings is not None:
             capture_warnings.extend(guard.warnings)
 

src/ha_mcp/dashboard_screenshot/theme_guard.py

@@ -1,12 +1,6 @@
 """Snapshot and restore the engine user's saved frontend theme (issue #1909).
 
-NOTE: This guard is currently disabled at its call site (``capture.py``)
-because upstream Puppet fixed the cold-render ``settheme`` dispatch that made
-the bracket necessary (#1991). The code here is retained unchanged so the
-bracket can be re-enabled by uncommenting the snapshot/restore calls in
-``capture.py`` if a future engine regression reintroduces the write.
-
-Stock Puppet dispatches Home Assistant's ``settheme`` event on every
+Puppet dispatches Home Assistant's ``settheme`` event on every
 cold-browser render — its ``dark`` query flag is presence-based, so "not
 requested" reaches the frontend as an explicit "light". Home Assistant
 persists that selection server-side per user (``frontend/set_user_data``,

tests/src/unit/test_dashboard_screenshot_theme_guard.py

@@ -5,13 +5,10 @@
 flipping that user's real sessions to light mode. The guard snapshots the
 saved theme before a capture batch and restores it afterwards.
 
-The capture-time bracket is currently DISABLED (#1991): upstream Puppet fixed
-the cold-render dispatch, so ``capture.py`` no longer calls the guard's
-snapshot/restore. The guard code itself is retained (and unchanged), so the
-unit tests below still cover its credential resolution and snapshot/restore
-semantics against a fake WebSocket client. ``TestCaptureBracketDisabled``
-asserts the bracket stays off — capture opens no guard sessions and adds no
-warnings — guarding against a silent re-enable.
+These cover the guard's credential resolution and snapshot/restore semantics
+against a fake WebSocket client, plus ``TestCaptureBracket``, which asserts
+``capture.py`` actually brackets each batch — guarding against the bracket
+being silently disabled again.
 """
 
 from __future__ import annotations
@@ -366,13 +363,18 @@ async def fake_resolve() -> EngineTarget:
     _ClobberingEngineClient.status_code = 200
 
 
-class TestCaptureBracketDisabled:
-    """The ThemeGuard bracket around capture_dashboard_images is DISABLED
-    (#1991). Upstream Puppet no longer clobbers the theme on cold renders, so
-    ha-mcp opens no snapshot/restore sessions and adds no guard warnings. These
-    guard against the bracket being silently re-enabled."""
+class TestCaptureBracket:
+    """The ThemeGuard bracket around ``capture_dashboard_images`` is ENABLED.
 
-    async def test_capture_opens_no_guard_sessions_and_leaves_theme(
+    Puppet dispatches a ``settheme`` event on render, which Home Assistant
+    persists onto the engine token user's profile (#1909), so every capture
+    batch snapshots the saved theme before and restores it after. The bracket
+    was once disabled on the premise that upstream Puppet had fixed the
+    dispatch; that fix (balloob/home-assistant-addons#89) is unreleased and is
+    scoped to the no-parameter case, so the bracket stays on.
+    """
+
+    async def test_capture_restores_theme_clobbered_by_the_render(
         self, monkeypatch: Any
     ) -> None:
         from ha_mcp.dashboard_screenshot import capture
@@ -386,19 +388,19 @@ async def test_capture_opens_no_guard_sessions_and_leaves_theme(
         )
 
         assert captures[0].data == _PNG
-        # No snapshot/restore WebSocket sessions were opened.
-        assert _FakeWsClient.instances == []
-        # The bracket did not run: the fake engine still simulates the old
-        # clobber, and ha-mcp no longer undoes it (real Puppet no longer
-        # causes it).
-        assert _FakeWsClient.user_data[THEME_USER_DATA_KEY] == _CLOBBERED_THEME
+        # Snapshot and restore each opened a session as the engine user.
+        assert [(ws.url, ws.token) for ws in _FakeWsClient.instances] == [
+            ("http://homeassistant:8123", "puppet-token"),
+            ("http://homeassistant:8123", "puppet-token"),
+        ]
+        # The fake engine clobbered the stored theme; the bracket undid it.
+        assert _FakeWsClient.user_data[THEME_USER_DATA_KEY] == _DARK_THEME
         assert capture_warnings == []
 
-    async def test_client_credential_fallback_opens_no_guard_sessions(
+    async def test_client_credential_fallback_restores_theme(
         self, monkeypatch: Any
     ) -> None:
-        """Sidecar/standalone mode: even with a usable client credential the
-        disabled bracket opens no guard sessions."""
+        """Sidecar/standalone mode: the client credential drives the bracket."""
         from ha_mcp.dashboard_screenshot import capture
 
         _patch_engine(monkeypatch, None)
@@ -409,11 +411,14 @@ async def test_client_credential_fallback_opens_no_guard_sessions(
         )
 
         assert captures[0].data == _PNG
-        assert _FakeWsClient.instances == []
+        assert [(ws.url, ws.token) for ws in _FakeWsClient.instances] == [
+            ("http://ha.local:8123", "own-token"),
+            ("http://ha.local:8123", "own-token"),
+        ]
+        assert _FakeWsClient.user_data[THEME_USER_DATA_KEY] == _DARK_THEME
 
-    async def test_capture_failure_raises_without_guard_warnings(
-        self, monkeypatch: Any
-    ) -> None:
+    async def test_capture_failure_still_restores_theme(self, monkeypatch: Any) -> None:
+        """A failed render must not leave the engine user's theme clobbered."""
         import json
 
         from ha_mcp.dashboard_screenshot import capture
@@ -425,8 +430,8 @@ async def test_capture_failure_raises_without_guard_warnings(
         with pytest.raises(ToolError) as exc_info:
             await capture.capture_dashboard_images("lovelace/0")
 
-        # No guard sessions opened, and no guard warning on the error payload.
-        assert _FakeWsClient.instances == []
+        # The restore in the ``finally`` ran before the error surfaced.
+        assert _FakeWsClient.user_data[THEME_USER_DATA_KEY] == _DARK_THEME
         payload = json.loads(str(exc_info.value))
         assert not any(
             "restoring it failed" in warning for warning in payload.get("warnings", [])