homeassistant-ai
diff --git a/‎src/ha_mcp/tools/backup.py‎
Lines changed: 122 additions & 15 deletions b/‎src/ha_mcp/tools/backup.py‎
Lines changed: 122 additions & 15 deletions
diff --git a/‎tests/src/unit/test_backup_agent_lookup.py‎
Lines changed: 42 additions & 15 deletions b/‎tests/src/unit/test_backup_agent_lookup.py‎
Lines changed: 42 additions & 15 deletions
@@ -52,6 +52,12 @@
 # happy-path wait noticeable.
 _BACKUP_MAX_WAIT_S = 300
 _BACKUP_POLL_INTERVAL_S = 2
+# The pre-restore safety backup is a *full* backup (include_database=True, plus
+# all add-ons on Supervised), unlike the fast create_backup path. A multi-GB
+# full backup on constrained hardware (~2.1 GB on a Pi 5, #1681) can run well
+# past the fast-backup window, so its completion poll gets a larger budget —
+# timing out here aborts the restore and the retry spawns another full backup.
+_SAFETY_BACKUP_MAX_WAIT_S = 1800
 # Clock-skew tolerance when filtering backup entries by date vs job-start.
 _BACKUP_DATE_FILTER_TOLERANCE_S = 5
 
@@ -540,17 +546,23 @@ async def _create_safety_backup(
     ws_client: HomeAssistantWebSocketClient,
     password: str | None,
     agent_id: str,
-) -> str | None:
-    """Create a pre-restore safety backup.
+) -> tuple[str | None, list[str]]:
+    """Create a pre-restore safety backup and wait for it to complete.
 
     ``agent_id`` is the local backup agent (Supervisor's ``hassio.local`` or
     Core's ``backup.local``) discovered by the caller.
 
-    Returns the safety backup ID, or None when password is None (backup intentionally
-    skipped). Raises ToolError if backup creation fails.
+    Blocks until the safety backup finishes. HA's ``backup/generate`` returns
+    on job initiation, not completion, so waiting here lets the caller issue
+    ``backup/restore`` afterwards without colliding with a still-running
+    backup. Returns ``(job_id, warnings)`` — ``job_id`` is None when password
+    is None (backup intentionally skipped); ``warnings`` carries any late-
+    completion notice from the poll so the caller can surface it before the
+    destructive restore. Raises ToolError if the backup fails to start or does
+    not complete in time.
     """
     if password is None:
-        return None
+        return None, []
 
     now = datetime.now()
     safety_backup_name = f"PreRestore_Safety_{now.strftime('%Y-%m-%d_%H:%M:%S')}"
@@ -579,8 +591,45 @@ async def _create_safety_backup(
         )
 
     safety_backup_id = safety_backup.get("result", {}).get("backup_job_id")
-    logger.info(f"Safety backup created: {safety_backup_id}")
-    return cast(str, safety_backup_id)
+    logger.info(f"Safety backup started: {safety_backup_id}, waiting for completion...")
+
+    # Wait for the safety backup to finish before returning. HA's
+    # backup/generate WS command returns as soon as the job is *initiated*,
+    # not when it completes, and the backup manager rejects any new operation
+    # while a backup is running ("Backup manager busy: create_backup"). If the
+    # caller issued backup/restore right after this returned, it would collide
+    # with the safety backup this same call just started – a self-induced
+    # deadlock (#1681). Reuse the same completion poll create_backup already
+    # uses rather than adding a second wait mechanism.
+    poll_result = await _poll_backup_completion(
+        ws_client,
+        safety_backup_name,
+        cast(str, safety_backup_id),
+        max_wait_seconds=_SAFETY_BACKUP_MAX_WAIT_S,
+        poll_interval=_BACKUP_POLL_INTERVAL_S,
+        agent_id=agent_id,
+    )
+    logger.info(f"Safety backup completed: {safety_backup_id}")
+    return cast(str, safety_backup_id), poll_result.get("warnings", [])
+
+
+def _backup_protected(entry: dict[str, Any]) -> bool | None:
+    """Whether a ``backup/info`` entry is encrypted.
+
+    ``protected`` is a per-agent field (AgentBackupStatus), not top-level. A
+    backup is encrypted as a whole, so any agent reporting it makes the backup
+    protected; returns None when the ``agents`` map is absent/malformed or no
+    agent reports the flag (unknown rather than "unprotected").
+    """
+    agents = entry.get("agents")
+    if not isinstance(agents, dict):
+        return None
+    flags = [
+        a.get("protected")
+        for a in agents.values()
+        if isinstance(a, dict) and a.get("protected") is not None
+    ]
+    return any(flags) if flags else None
 
 
 async def restore_backup(
@@ -594,7 +643,12 @@ async def restore_backup(
     Args:
         client: Home Assistant REST client
         backup_id: Backup ID to restore
-        restore_database: Whether to restore database (historical data)
+        restore_database: Whether to restore database (historical data).
+            On Supervised installs this is reconciled against the target
+            backup's ``database_included`` flag, which Supervisor requires to
+            match when restoring Home Assistant; a warning is returned if the
+            value is overridden. On Core installs the caller's value is honoured
+            as-is (Core enforces no such match).
 
     Returns:
         Dictionary with restore result including safety_backup_id, status, etc.
@@ -627,9 +681,9 @@ async def restore_backup(
             )
 
         backups = backup_info.get("result", {}).get("backups", [])
-        backup_exists = any(b.get("backup_id") == backup_id for b in backups)
+        matched = next((b for b in backups if b.get("backup_id") == backup_id), None)
 
-        if not backup_exists:
+        if matched is None:
             raise_tool_error(
                 create_error_response(
                     ErrorCode.RESOURCE_NOT_FOUND,
@@ -655,17 +709,57 @@ async def restore_backup(
             logger.warning("No default password - proceeding without safety backup")
             password = None
 
-        safety_backup_id = await _create_safety_backup(ws_client, password, local_agent)
+        safety_backup_id, safety_warnings = await _create_safety_backup(
+            ws_client, password, local_agent
+        )
+
+        # `backup/info` returns ManagerBackup entries: `database_included` is a
+        # top-level field (inherited from BaseBackup), but `protected` is NOT —
+        # it lives per-agent under the entry's `agents` map (AgentBackupStatus),
+        # so a top-level `matched.get("protected")` is always None against real
+        # HA. Derive it from the agents map (shared with _summarize_backup).
+        target_protected = _backup_protected(matched)
+
+        # Reconcile restore_database with the target only on Supervised. HA's
+        # Supervisor raises "Restore database must match backup" when
+        # restore_homeassistant is set and restore_database != the backup's
+        # database_included flag (hassio/backup.py). HA Core's restore path
+        # (CoreBackupReaderWriter) has no such constraint and writes the caller's
+        # value verbatim, so overriding it there would silently discard an
+        # explicit request for no HA-side reason. Fall back to the caller's
+        # request when not Supervised or when the field is absent. Surface a
+        # warning when the derived value overrides what the caller asked for so
+        # the override isn't silent on a destructive op.
+        is_supervised = local_agent == "hassio.local"
+        target_database_included = matched.get("database_included")
+        if is_supervised and target_database_included is not None:
+            effective_restore_database = target_database_included
+        else:
+            effective_restore_database = restore_database
 
         # Perform restore
-        restore_params = {
+        restore_params: dict[str, Any] = {
             "backup_id": backup_id,
             "agent_id": local_agent,
-            "restore_database": restore_database,
+            "restore_database": effective_restore_database,
             "restore_homeassistant": True,
             "restore_addons": [],  # Restore all addons from backup
             "restore_folders": [],  # Restore all folders from backup
         }
+        # Forward the default backup password ONLY for a protected (encrypted)
+        # target. `password` here is HA's default create_backup.password, which
+        # is independent of whether *this* backup is encrypted: HA validates the
+        # password against the target unconditionally and rejects a password on
+        # an unprotected backup ("Invalid password for backup" →
+        # IncorrectPasswordError). Gate on the target's per-agent `protected`
+        # flag (read above) so an unprotected snapshot still restores on a
+        # default-password instance. HA's backup/restore schema types `password`
+        # as `str` (not `str | None`), so only include the key when we actually
+        # forward one — passing None would fail voluptuous validation. Without
+        # this a protected backup cannot be restored even though the HA UI,
+        # which applies the stored key, succeeds (#1681).
+        if password is not None and target_protected:
+            restore_params["password"] = password
 
         result = await ws_client.send_command("backup/restore", **restore_params)
 
@@ -678,6 +772,18 @@ async def restore_backup(
             warnings = [
                 "Home Assistant is restarting. Connection will be temporarily lost."
             ]
+            # Surface any late-completion notice from the safety-backup poll —
+            # a slow backup subsystem right before a destructive restore is
+            # exactly the signal a caller wants to see.
+            warnings.extend(safety_warnings)
+            if effective_restore_database != restore_database:
+                warnings.append(
+                    "restore_database was adjusted to "
+                    f"{effective_restore_database} to match the target backup "
+                    "(Home Assistant's Supervisor requires it to match the "
+                    "backup's database_included flag when restoring Home "
+                    "Assistant)."
+                )
             if safety_backup_id is None:
                 warnings.append(
                     "No safety backup was created (the default backup "
@@ -697,7 +803,7 @@ async def restore_backup(
                 "backup_id": backup_id,
                 "status": "Restore initiated - Home Assistant will restart",
                 "safety_backup_id": safety_backup_id,
-                "restore_database": restore_database,
+                "restore_database": effective_restore_database,
                 "warnings": warnings,
                 "note": note,
             }
@@ -761,7 +867,8 @@ def _summarize_backup(entry: dict[str, Any]) -> dict[str, Any]:
         "name": entry.get("name"),
         "date": entry.get("date"),
         "size_bytes": size_bytes,
-        "protected": entry.get("protected"),
+        # per-agent field, derived from the agents map (see _backup_protected)
+        "protected": _backup_protected(entry),
         "database_included": entry.get("database_included"),
         "homeassistant_included": entry.get("homeassistant_included"),
         "homeassistant_version": entry.get("homeassistant_version"),
 
@@ -11,13 +11,41 @@
 from fastmcp.exceptions import ToolError
 
 from ha_mcp.tools.backup import (
+    _backup_protected,
     _get_local_backup_agent_id,
     _summarize_backup,
     list_backups,
     restore_backup,
 )
 
 
+class TestBackupProtected:
+    """`_backup_protected` derives the encrypted flag from the per-agent map."""
+
+    def test_any_agent_protected_makes_it_protected(self):
+        entry = {
+            "agents": {
+                "backup.local": {"protected": False},
+                "google_drive.cloud": {"protected": True},
+            }
+        }
+        assert _backup_protected(entry) is True
+
+    def test_all_agents_unprotected(self):
+        entry = {"agents": {"backup.local": {"protected": False}}}
+        assert _backup_protected(entry) is False
+
+    def test_missing_agents_is_unknown(self):
+        assert _backup_protected({}) is None
+
+    def test_no_agent_reports_protected_is_unknown(self):
+        # agents present but none carry the field → unknown, not "unprotected".
+        assert _backup_protected({"agents": {"backup.local": {"size": 1}}}) is None
+
+    def test_non_dict_agents_is_unknown(self):
+        assert _backup_protected({"agents": ["unexpected"]}) is None
+
+
 def _ws_client(agents_payload: dict) -> AsyncMock:
     """Build a mock WS client whose `send_command("backup/agents/info")` returns the given payload."""
     ws = AsyncMock()
@@ -157,15 +185,6 @@ async def test_success_returns_top_level_warnings_list(self):
             },
             # _create_safety_backup → backup/generate
             {"success": True, "result": {"backup_job_id": "safety_job_1"}},
-            # _create_safety_backup polling → backup/info loop
-            {
-                "success": True,
-                "result": {
-                    "backups": [
-                        {"name": "Pre_Restore_Safety", "backup_id": "safety_xyz"}
-                    ]
-                },
-            },
             # backup/restore — the actual restore call
             {"success": True},
         ]
@@ -175,9 +194,17 @@ async def test_success_returns_top_level_warnings_list(self):
         client.token = "token"
         client.verify_ssl = False
 
-        with patch(
-            "ha_mcp.tools.backup.get_connected_ws_client",
-            new=AsyncMock(return_value=(ws, None)),
+        # The safety-backup completion poll is exercised by test_backup_restore;
+        # stub it here so this test stays focused on the warnings contract.
+        with (
+            patch(
+                "ha_mcp.tools.backup.get_connected_ws_client",
+                new=AsyncMock(return_value=(ws, None)),
+            ),
+            patch(
+                "ha_mcp.tools.backup._poll_backup_completion",
+                new=AsyncMock(return_value={"success": True}),
+            ),
         ):
             result = await restore_backup(client, "abc123")
 
@@ -200,14 +227,14 @@ def test_projects_core_fields_and_largest_agent_size(self):
             "backup_id": "abc",
             "name": "Nightly",
             "date": "2026-06-14T02:00:00+00:00",
-            "protected": True,
             "database_included": False,
             "homeassistant_included": True,
             "homeassistant_version": "2026.6.0",
             "with_automatic_settings": True,
+            # `protected` is a per-agent field (AgentBackupStatus), not top-level.
             "agents": {
-                "backup.local": {"size": 100},
-                "google_drive.cloud": {"size": 250},
+                "backup.local": {"size": 100, "protected": True},
+                "google_drive.cloud": {"size": 250, "protected": True},
             },
         }
         out = _summarize_backup(entry)