ha-mcp-dev 7.13.0.dev1975__tar.gz → 7.13.0.dev1980__tar.gz

{ha_mcp_dev-7.13.0.dev1975/src/ha_mcp_dev.egg-info → ha_mcp_dev-7.13.0.dev1980}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ha-mcp-dev
-Version: 7.13.0.dev1975
+Version: 7.13.0.dev1980
 Summary: Home Assistant MCP Server - Complete control of Home Assistant through MCP
 Author-email: Julien <github@qc-h.net>
 License: MIT
@@ -118,6 +118,7 @@ These run the server outside Home Assistant — useful for **Container** / **Cor
 - **Docker (HTTP server):** run `ghcr.io/homeassistant-ai/ha-mcp` in HTTP mode, pointed at your Home Assistant URL and a long-lived token, and connect your client to its secret URL. See the [Setup Wizard](https://homeassistant-ai.github.io/ha-mcp/setup/) for the full command and per-client config.
 - **PyPI / uvx (HTTP server):** run the published `ha-mcp` package with `uvx ha-mcp@latest` (or pip) as a streamable-HTTP server the same way. Details in the [Setup Wizard](https://homeassistant-ai.github.io/ha-mcp/setup/).
 - **Local stdio (not recommended):** runs ha-mcp on your own machine over stdio. The one-command installers in the **Demo server** section below use this path; the [Setup Wizard](https://homeassistant-ai.github.io/ha-mcp/setup/) covers connecting it to your own Home Assistant.
+- **OIDC authentication:** gate remote access behind an external identity provider (Authentik, Keycloak, Auth0, Google, etc.) instead of a secret URL — all authenticated users share the server's Home Assistant credentials. See [OIDC Mode](docs/oidc.md).
 
   > ⚠️ **stdio has known transport issues.** The stdio transport has connection problems that streamable HTTP does not ([#1713](https://github.com/homeassistant-ai/ha-mcp/issues/1713)). It is recommended only for demo/testing tinkering — for a real setup, use the custom component or an HTTP method above.
 

{ha_mcp_dev-7.13.0.dev1975 → ha_mcp_dev-7.13.0.dev1980}/README.md

@@ -85,6 +85,7 @@ These run the server outside Home Assistant — useful for **Container** / **Cor
 - **Docker (HTTP server):** run `ghcr.io/homeassistant-ai/ha-mcp` in HTTP mode, pointed at your Home Assistant URL and a long-lived token, and connect your client to its secret URL. See the [Setup Wizard](https://homeassistant-ai.github.io/ha-mcp/setup/) for the full command and per-client config.
 - **PyPI / uvx (HTTP server):** run the published `ha-mcp` package with `uvx ha-mcp@latest` (or pip) as a streamable-HTTP server the same way. Details in the [Setup Wizard](https://homeassistant-ai.github.io/ha-mcp/setup/).
 - **Local stdio (not recommended):** runs ha-mcp on your own machine over stdio. The one-command installers in the **Demo server** section below use this path; the [Setup Wizard](https://homeassistant-ai.github.io/ha-mcp/setup/) covers connecting it to your own Home Assistant.
+- **OIDC authentication:** gate remote access behind an external identity provider (Authentik, Keycloak, Auth0, Google, etc.) instead of a secret URL — all authenticated users share the server's Home Assistant credentials. See [OIDC Mode](docs/oidc.md).
 
   > ⚠️ **stdio has known transport issues.** The stdio transport has connection problems that streamable HTTP does not ([#1713](https://github.com/homeassistant-ai/ha-mcp/issues/1713)). It is recommended only for demo/testing tinkering — for a real setup, use the custom component or an HTTP method above.
 

{ha_mcp_dev-7.13.0.dev1975 → ha_mcp_dev-7.13.0.dev1980}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "ha-mcp-dev"
-version = "7.13.0.dev1975"
+version = "7.13.0.dev1980"
 description = "Home Assistant MCP Server - Complete control of Home Assistant through MCP"
 readme = "README.md"
 # Upper bound widened to <3.15 so the package installs on Home Assistant's
@@ -59,6 +59,7 @@ hamcp-test-env = "tests.test_env_manager:main"
 ha-mcp-web = "ha_mcp.__main__:main_web"
 ha-mcp-sse = "ha_mcp.__main__:main_sse"
 ha-mcp-oauth = "ha_mcp.__main__:main_oauth"
+ha-mcp-oidc = "ha_mcp.__main__:main_oidc"
 
 [tool.setuptools]
 package-dir = {"" = "src", "tests" = "tests"}

{ha_mcp_dev-7.13.0.dev1975 → ha_mcp_dev-7.13.0.dev1980}/src/ha_mcp/__main__.py

@@ -979,6 +979,39 @@ def _healthz_enabled() -> bool:
     return os.getenv("MCP_HEALTHZ", "").strip().lower() in ("1", "true", "yes", "on")
 
 
+def _oidc_verify_id_token_enabled() -> bool:
+    """True when OIDC_VERIFY_ID_TOKEN opts in to ID-token verification.
+
+    Off by default: FastMCP's OIDCProxy verifies the access token as a JWT,
+    which works for providers like Authentik and Keycloak. Providers that
+    issue opaque access tokens (Google always; Auth0 without an API
+    audience) need ``verify_id_token=True`` so FastMCP verifies the ID token
+    instead.
+    """
+    return os.getenv("OIDC_VERIFY_ID_TOKEN", "").strip().lower() in (
+        "1",
+        "true",
+        "yes",
+        "on",
+    )
+
+
+def _oidc_allowed_client_redirect_uris() -> list[str] | None:
+    """Parse OIDC_ALLOWED_CLIENT_REDIRECT_URIS into a list, or None if unset.
+
+    Comma-separated list of redirect URI patterns FastMCP's OIDCProxy will
+    accept from dynamically-registered clients. With no allow-list set,
+    FastMCP still validates a client's redirect URI against that same
+    client's own registered redirect URIs — the exposure is that open DCR
+    lets an attacker register their own client with their own redirect URI;
+    internet-facing deployments should set this to constrain what any
+    dynamically-registered client may register in the first place.
+    """
+    raw = os.getenv("OIDC_ALLOWED_CLIENT_REDIRECT_URIS", "")
+    uris = [uri.strip() for uri in raw.split(",") if uri.strip()]
+    return uris or None
+
+
 def register_browser_landing(
     mcp_instance: "FastMCP | _DeferredMCP",
     path: str,
@@ -1225,5 +1258,192 @@ async def _run_oauth_server(
     )
 
 
+def main_oidc() -> None:
+    """Run server with OIDC authentication over HTTP.
+
+    This mode enables authentication via an external OIDC provider
+    (Authentik, Keycloak, Auth0, Google, etc.). All authenticated users
+    share the same Home Assistant instance via the configured credentials.
+
+    Unlike OAuth mode which collects per-user HA credentials via a consent form,
+    OIDC mode is purely an access gate — users authenticate through the OIDC provider,
+    then all requests use the shared HA credentials from environment variables.
+
+    Environment:
+    - OIDC_CONFIG_URL (required): OIDC discovery URL (.well-known/openid-configuration)
+    - OIDC_CLIENT_ID (required): OAuth client ID registered with your OIDC provider
+    - OIDC_CLIENT_SECRET (required): OAuth client secret from your OIDC provider
+    - MCP_BASE_URL (required): Public HTTPS URL where this server is accessible
+    - HOMEASSISTANT_URL (required): Home Assistant instance URL
+    - HOMEASSISTANT_TOKEN (required): Home Assistant long-lived access token or supervisor token
+    - OIDC_JWT_SIGNING_KEY (optional): Secret key for signing FastMCP JWTs. Sessions
+      persist across restarts by default (the key is derived from OIDC_CLIENT_SECRET);
+      set this to decouple the signing key from the client secret. Generate with:
+      python -c "import secrets; print(secrets.token_urlsafe(32))"
+    - OIDC_ALLOWED_CLIENT_REDIRECT_URIS (optional): Comma-separated list of redirect URI
+      patterns accepted from dynamically-registered clients. Strongly recommended for
+      internet-facing deployments.
+    - OIDC_VERIFY_ID_TOKEN (optional, default: false): Set true for providers that issue
+      opaque access tokens (e.g. Google, or Auth0 without an API audience).
+    - OIDC_AUDIENCE (optional): Expected `aud` claim for IdP-issued access tokens.
+      Without it (and with OIDC_VERIFY_ID_TOKEN off), FastMCP's JWT verifier checks
+      issuer, signature, and expiry but not audience. With OIDC_VERIFY_ID_TOKEN=true,
+      verification pins `aud` to OIDC_CLIENT_ID instead and this value is only
+      forwarded to the IdP's authorize/token endpoints (it is what makes Auth0
+      issue JWT rather than opaque access tokens).
+    - MCP_HOST (optional, default: "0.0.0.0"; set 127.0.0.1 to restrict to loopback
+      behind a same-host reverse proxy)
+    - MCP_PORT (optional, default: 8086)
+    - MCP_SECRET_PATH (optional, default: "/mcp")
+    - LOG_LEVEL (optional, default: INFO)
+    """
+    # Configure logging for OIDC mode (force=True needed since logging may already be configured)
+    log_level = os.getenv("LOG_LEVEL", "INFO").upper()
+    _setup_logging(log_level, force=True)
+    for logger_name in ["ha_mcp", "ha_mcp.auth", "ha_mcp.auth.provider"]:
+        logging.getLogger(logger_name).setLevel(getattr(logging, log_level))
+    logger.info(f"OIDC mode logging configured at {log_level} level")
+    _log_startup_version()
+
+    # Validate HA credentials (OIDC mode needs them — unlike OAuth mode)
+    from ha_mcp.config import get_settings
+
+    settings = get_settings()
+    _validate_standard_credentials(settings)
+
+    # Validate OIDC-specific env vars
+    oidc_config_url = os.getenv("OIDC_CONFIG_URL")
+    oidc_client_id = os.getenv("OIDC_CLIENT_ID")
+    oidc_client_secret = os.getenv("OIDC_CLIENT_SECRET")
+    base_url = os.getenv("MCP_BASE_URL")
+
+    missing = []
+    if not oidc_config_url:
+        missing.append("OIDC_CONFIG_URL")
+    if not oidc_client_id:
+        missing.append("OIDC_CLIENT_ID")
+    if not oidc_client_secret:
+        missing.append("OIDC_CLIENT_SECRET")
+    if not base_url:
+        missing.append("MCP_BASE_URL")
+
+    if missing:
+        logger.error(
+            f"Missing required environment variables for OIDC mode: {', '.join(missing)}"
+        )
+        logger.error(
+            "OIDC mode requires an external OIDC provider (Authentik, Keycloak, Auth0, etc.)"
+        )
+        sys.exit(1)
+
+    assert oidc_config_url is not None
+    assert oidc_client_id is not None
+    assert oidc_client_secret is not None
+    assert base_url is not None
+
+    host, port, path = _get_http_runtime(default_port=8086)
+
+    _run_entrypoint(
+        _run_oidc_server(
+            config_url=oidc_config_url,
+            client_id=oidc_client_id,
+            client_secret=oidc_client_secret,
+            base_url=base_url,
+            host=host,
+            port=port,
+            path=path,
+        ),
+        "OIDC server",
+    )
+
+
+async def _run_oidc_server(
+    config_url: str,
+    client_id: str,
+    client_secret: str,
+    base_url: str,
+    host: str,
+    port: int,
+    path: str,
+) -> None:
+    """Run the OIDC-authenticated MCP server.
+
+    Unlike OAuth mode which uses OAuthProxyClient for per-user credential routing,
+    OIDC mode uses the standard HomeAssistantClient with shared credentials.
+    OIDC is purely an access gate — all authenticated users share the same HA instance.
+
+    Args:
+        config_url: OIDC discovery URL (.well-known/openid-configuration)
+        client_id: OAuth client ID from the OIDC provider
+        client_secret: OAuth client secret from the OIDC provider
+        base_url: Public HTTPS URL where this server is accessible
+        host: Bind host (typically 0.0.0.0; override via MCP_HOST)
+        port: Port to listen on
+        path: MCP endpoint path
+    """
+    from fastmcp.server.auth.oidc_proxy import OIDCProxy
+
+    from ha_mcp.server import HomeAssistantSmartMCPServer
+    from ha_mcp.transport_security import ensure_host_origin_guard_default_off
+
+    # OIDC mode is deployed behind a public reverse proxy or LAN hostname;
+    # disable FastMCP's Host/Origin rebinding guard so proxied requests are
+    # not rejected with 421/403 before OIDC auth runs.
+    ensure_host_origin_guard_default_off()
+
+    proxy_kwargs: dict[str, Any] = {
+        "config_url": config_url,
+        "client_id": client_id,
+        "client_secret": client_secret,
+        "base_url": base_url,
+        # "external" tells FastMCP that consent is handled by the upstream
+        # IdP (Authentik, Keycloak, etc.) -- unlike `False`, this does not
+        # log a security warning at startup that consent is disabled.
+        "require_authorization_consent": "external",
+        # Preserve `or None`: an empty-but-set env var must not bypass
+        # FastMCP's derive-from-client-secret default for jwt_signing_key.
+        "jwt_signing_key": os.getenv("OIDC_JWT_SIGNING_KEY") or None,
+    }
+
+    allowed_redirect_uris = _oidc_allowed_client_redirect_uris()
+    if allowed_redirect_uris:
+        proxy_kwargs["allowed_client_redirect_uris"] = allowed_redirect_uris
+    else:
+        logger.warning(
+            "OIDC_ALLOWED_CLIENT_REDIRECT_URIS is not set: any dynamically-registered "
+            "(DCR) client's own redirect URIs are accepted. Internet-facing "
+            "deployments should set it."
+        )
+
+    if _oidc_verify_id_token_enabled():
+        proxy_kwargs["verify_id_token"] = True
+
+    audience = os.getenv("OIDC_AUDIENCE", "").strip()
+    if audience:
+        proxy_kwargs["audience"] = audience
+
+    # Create OIDC auth provider — auto-discovers endpoints from config_url
+    auth = OIDCProxy(**proxy_kwargs)
+
+    # Standard server with shared credentials (no proxy client needed)
+    global _server
+    _server = HomeAssistantSmartMCPServer()
+    mcp_instance = _server.mcp
+    mcp_instance.auth = auth
+
+    logger.info("Server created with OIDC authentication")
+    # Unlike standard/OAuth HTTP modes, OIDC mode intentionally registers no
+    # browser landing page and no settings-UI sidecar: an unauthenticated
+    # browser GET gets a bare 405 rather than a friendly page, and there is
+    # no settings UI surface here.
+    if _healthz_enabled():
+        _register_healthz_route(mcp_instance)
+    logger.info(f"Starting OIDC-enabled MCP server at {base_url}{path}")
+
+    await _run_with_shutdown(
+        mcp_instance.run_async(**_http_run_kwargs("http", host, port, path))
+    )
+
+
 if __name__ == "__main__":
     main()

{ha_mcp_dev-7.13.0.dev1975 → ha_mcp_dev-7.13.0.dev1980}/src/ha_mcp/config.py

@@ -836,7 +836,7 @@ ADVANCED_SETTINGS_FIELDS: tuple[AdvancedField, ...] = (
     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(

{ha_mcp_dev-7.13.0.dev1975 → ha_mcp_dev-7.13.0.dev1980}/src/ha_mcp/dashboard_screenshot/__init__.py

@@ -19,7 +19,8 @@ from .capture import (
     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",
@@ -27,6 +28,9 @@ __all__ = [
     "DEFAULT_WAIT_MS",
     "DEFAULT_WIDTH",
     "DashboardImageCapture",
+    "EngineCredential",
+    "EngineTarget",
+    "ThemeGuard",
     "capture_dashboard_images",
-    "resolve_engine_url",
+    "resolve_engine",
 ]

{ha_mcp_dev-7.13.0.dev1975 → ha_mcp_dev-7.13.0.dev1980}/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 fastmcp.exceptions import ToolError
 
 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,6 +845,8 @@ 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.
 
@@ -796,6 +854,12 @@ async def capture_dashboard_images(
     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)