From c1757c8ca325a3a797e07c6c305495e4c97d01bf Mon Sep 17 00:00:00 2001
From: Filipe Casal <fcasal@users.noreply.github.com>
Date: Tue, 8 Sep 2026 13:09:04 +0100
Subject: [PATCH 1/2] Fix wake recovery after snapshot restore timeout

---
 README.md          |   3 +
 dropkit/api.py     |  15 ++--
 dropkit/main.py    | 106 +++++++++++++++++-----
 tests/test_wake.py | 216 +++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 311 insertions(+), 29 deletions(-)
 create mode 100644 tests/test_wake.py

diff --git a/README.md b/README.md
index 68b7c68..09a59ef 100644
--- a/README.md
+++ b/README.md
@@ -178,6 +178,9 @@ dropkit destroy my-droplet
 3. Snapshots are tagged with `owner:<username>` and `size:<size-slug>` for tracking
 4. After waking, you're prompted to delete the snapshot (default: yes)
 
+Wake waits up to 15 minutes by default. Use `dropkit wake my-droplet --timeout 1800`
+to wait up to 30 minutes.
+
 **Note:** Snapshots are billed at $0.06/GB/month, which is typically much cheaper than keeping a droplet running.
 
 ### Cloud-Init Customization
diff --git a/dropkit/api.py b/dropkit/api.py
index 47e09b3..d3c93a6 100644
--- a/dropkit/api.py
+++ b/dropkit/api.py
@@ -68,6 +68,10 @@ def __init__(self, message: str, status_code: int | None = None):
         super().__init__(message)
 
 
+class DropletWaitTimeoutError(DigitalOceanAPIError):
+    """The local wait expired; droplet creation may still be in progress."""
+
+
 PROTECTED_TAGS = {"owner", "firewall"}
 
 
@@ -405,12 +409,13 @@ def wait_for_droplet_active(
 
         Raises:
             ValueError: If droplet_id is not positive
-            DigitalOceanAPIError: If timeout is reached or droplet enters error state
+            DropletWaitTimeoutError: If the local wait expires
+            DigitalOceanAPIError: If the droplet enters an error state or an API request fails
         """
         import time
 
         self._validate_positive_int(droplet_id, "droplet_id")
-        start_time = time.time()
+        start_time = time.monotonic()
 
         while True:
             droplet = self.get_droplet(droplet_id)
@@ -423,9 +428,9 @@ def wait_for_droplet_active(
                     f"Droplet entered error state: {droplet.get('name', droplet_id)}"
                 )
 
-            elapsed = time.time() - start_time
-            if elapsed > timeout:
-                raise DigitalOceanAPIError(
+            elapsed = time.monotonic() - start_time
+            if elapsed >= timeout:
+                raise DropletWaitTimeoutError(
                     f"Timeout waiting for droplet to become active (waited {elapsed:.0f}s)"
                 )
 
diff --git a/dropkit/main.py b/dropkit/main.py
index 7b31fa1..b34cad5 100644
--- a/dropkit/main.py
+++ b/dropkit/main.py
@@ -16,7 +16,7 @@
 from rich.prompt import Confirm, Prompt
 from rich.table import Table
 
-from dropkit.api import DigitalOceanAPI, DigitalOceanAPIError
+from dropkit.api import DigitalOceanAPI, DigitalOceanAPIError, DropletWaitTimeoutError
 from dropkit.cloudinit import render_cloud_init
 from dropkit.config import DEFAULT_IMAGE, DEFAULT_REGION, DEFAULT_SIZE, Config, DropkitConfig
 from dropkit.lock import requires_lock
@@ -4036,12 +4036,16 @@ def wake(
     droplet_name: str = typer.Argument(
         ..., autocompletion=complete_snapshot_name, help="Name of the hibernated droplet to restore"
     ),
+    timeout: int = typer.Option(
+        900, "--timeout", min=1, help="Seconds to wait for restoration (default: 15 minutes)"
+    ),
     no_tailscale: bool = typer.Option(False, "--no-tailscale", help="Skip Tailscale VPN re-setup"),
 ):
     """
     Wake a hibernated droplet (restore from snapshot).
 
     This will create a new droplet from the hibernated snapshot.
+    Run again after a timeout to resume setup of the same restored droplet.
     After successful restoration, you'll be prompted to delete the snapshot.
 
     If the original droplet had Tailscale lockdown enabled, this command will
@@ -4064,10 +4068,6 @@ def wake(
 
         # Check if a droplet with this name already exists
         existing_droplet, _ = find_user_droplet(api, droplet_name)
-        if existing_droplet:
-            console.print(f"[red]Error: A droplet named '{droplet_name}' already exists.[/red]")
-            console.print("[dim]Destroy or rename the existing droplet first.[/dim]")
-            raise typer.Exit(1)
 
         # Find the hibernated snapshot
         snapshot_name = get_snapshot_name(droplet_name)
@@ -4079,6 +4079,13 @@ def wake(
         if not snapshot:
             console.print(f"[red]Error: No hibernated snapshot found for '{droplet_name}'[/red]")
             console.print(f"[dim]Expected snapshot name: {snapshot_name}[/dim]")
+            if existing_droplet:
+                console.print(
+                    f"[dim]The droplet already exists. Check it with: dropkit info {droplet_name}[/dim]"
+                )
+                console.print(
+                    f"[dim]To configure SSH, run: dropkit config-ssh {droplet_name}[/dim]"
+                )
             raise typer.Exit(1)
 
         snapshot_id_str = snapshot.get("id")
@@ -4087,6 +4094,23 @@ def wake(
             raise typer.Exit(1)
         snapshot_id = int(snapshot_id_str)  # API returns string, convert to int
 
+        if existing_droplet:
+            # A matching name alone is insufficient: only resume a restore of this snapshot.
+            source_image_id = existing_droplet.get("image", {}).get("id")
+            if str(source_image_id) != str(snapshot_id):
+                console.print(
+                    f"[red]Error: A droplet named '{droplet_name}' already exists, "
+                    "but was not restored from this snapshot.[/red]"
+                )
+                console.print(f"[dim]Inspect it with: dropkit info {droplet_name}[/dim]")
+                raise typer.Exit(1)
+            status = existing_droplet.get("status")
+            if status not in {"new", "active"}:
+                console.print(f"[red]Cannot resume wake: droplet status is '{status}'.[/red]")
+                if status == "off":
+                    console.print(f"[dim]Power it on first: dropkit on {droplet_name}[/dim]")
+                raise typer.Exit(1)
+
         # Get snapshot details
         size_gb = snapshot.get("size_gigabytes", 0)
         regions = snapshot.get("regions", [])
@@ -4120,33 +4144,60 @@ def wake(
         )
         console.print()
 
-        # Create droplet from snapshot
-        console.print(f"[dim]Creating droplet '{droplet_name}' from snapshot...[/dim]")
-
-        # Build tags for new droplet
-        tags_list = build_droplet_tags(username, list(config.defaults.extra_tags))
-
-        droplet = api.create_droplet_from_snapshot(
-            name=droplet_name,
-            region=original_region,
-            size=original_size,
-            snapshot_id=snapshot_id,
-            tags=tags_list,
-            ssh_keys=config.cloudinit.ssh_key_ids,
-        )
+        if existing_droplet:
+            droplet = existing_droplet
+            console.print(f"[dim]Resuming wake for existing droplet '{droplet_name}'...[/dim]")
+        else:
+            # Create droplet from snapshot
+            console.print(f"[dim]Creating droplet '{droplet_name}' from snapshot...[/dim]")
+
+            # Build tags for new droplet
+            tags_list = build_droplet_tags(username, list(config.defaults.extra_tags))
+
+            droplet = api.create_droplet_from_snapshot(
+                name=droplet_name,
+                region=original_region,
+                size=original_size,
+                snapshot_id=snapshot_id,
+                tags=tags_list,
+                ssh_keys=config.cloudinit.ssh_key_ids,
+            )
 
         droplet_id = droplet.get("id")
         if not droplet_id:
             console.print("[red]Error: Failed to get droplet ID from API response[/red]")
             raise typer.Exit(1)
 
-        console.print(f"[green]✓[/green] Droplet created (ID: [cyan]{droplet_id}[/cyan])")
+        console.print(f"[green]✓[/green] Droplet ID: [cyan]{droplet_id}[/cyan]")
 
         # Wait for droplet to become active
         console.print("[dim]Waiting for droplet to become active...[/dim]")
 
-        with console.status("[cyan]Waiting...[/cyan]"):
-            active_droplet = api.wait_for_droplet_active(droplet_id)
+        try:
+            with console.status("[cyan]Waiting...[/cyan]"):
+                active_droplet = api.wait_for_droplet_active(droplet_id, timeout=timeout)
+        except DropletWaitTimeoutError:
+            console.print(
+                f"[yellow]Stopped waiting after {timeout}s. DigitalOcean may still be "
+                f"restoring droplet '{droplet_name}' (ID: {droplet_id}).[/yellow]"
+            )
+            console.print("[dim]The droplet and snapshot have been kept.[/dim]")
+            console.print(f"[dim]Check status: dropkit info {droplet_name}[/dim]")
+            console.print(
+                f"[dim]Resume setup: dropkit wake {droplet_name} --timeout {timeout}"
+                f"{' --no-tailscale' if no_tailscale else ''}[/dim]"
+            )
+            raise typer.Exit(1)
+        except DigitalOceanAPIError:
+            console.print(
+                f"[yellow]Could not check readiness of droplet {droplet_id}. "
+                "The droplet and snapshot have been kept.[/yellow]"
+            )
+            console.print(
+                f"[dim]Check status with dropkit info {droplet_name}, "
+                "then rerun wake to resume setup.[/dim]"
+            )
+            raise
 
         # Get IP address
         networks = active_droplet.get("networks", {})
@@ -4162,7 +4213,9 @@ def wake(
             console.print(f"[green]✓[/green] Droplet is active (IP: [cyan]{ip_address}[/cyan])")
         else:
             console.print("[green]✓[/green] Droplet is active")
-            console.print("[yellow]⚠[/yellow] Could not determine IP address")
+            console.print("[red]Could not determine public IP address; snapshot kept.[/red]")
+            console.print(f"[dim]Retry setup: dropkit wake {droplet_name}[/dim]")
+            raise typer.Exit(1)
 
         # Add SSH config entry
         if ip_address and config.ssh.auto_update:
@@ -4178,7 +4231,12 @@ def wake(
                 )
                 console.print("[green]✓[/green] SSH config updated")
             except Exception as e:
-                console.print(f"[yellow]⚠[/yellow] Could not update SSH config: {e}")
+                console.print(f"[red]Could not update SSH config: {e}[/red]")
+                console.print(
+                    f"[dim]Snapshot kept. Fix the SSH configuration error, "
+                    f"then retry: dropkit wake {droplet_name}[/dim]"
+                )
+                raise typer.Exit(1)
 
         # Handle Tailscale re-setup if the original droplet had Tailscale lockdown
         if was_tailscale_locked and ip_address:
diff --git a/tests/test_wake.py b/tests/test_wake.py
new file mode 100644
index 0000000..63cf791
--- /dev/null
+++ b/tests/test_wake.py
@@ -0,0 +1,216 @@
+"""Regression tests for interrupted snapshot restores."""
+
+from contextlib import ExitStack
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+from typer.testing import CliRunner
+
+from dropkit.api import DigitalOceanAPI, DigitalOceanAPIError, DropletWaitTimeoutError
+from dropkit.main import app
+
+runner = CliRunner()
+
+
+@pytest.fixture
+def wake_env(tmp_path):
+    """Isolate all cloud, SSH, VPN, and lock effects."""
+    api = MagicMock()
+    api.get_username.return_value = "testuser"
+    snapshot = {
+        "id": "456",
+        "regions": ["lon1"],
+        "tags": ["owner:testuser", "size:s-8vcpu-16gb"],
+    }
+    droplet = {
+        "id": 123,
+        "name": "test-droplet",
+        "image": {"id": 456},
+        "status": "active",
+        "networks": {"v4": [{"type": "public", "ip_address": "192.0.2.1"}]},
+    }
+    api.get_snapshot_by_name.return_value = snapshot
+    api.create_droplet_from_snapshot.return_value = droplet
+    api.wait_for_droplet_active.return_value = droplet
+    config = SimpleNamespace(
+        defaults=SimpleNamespace(extra_tags=[], size="s-8vcpu-16gb"),
+        cloudinit=SimpleNamespace(ssh_key_ids=[789]),
+        ssh=SimpleNamespace(
+            auto_update=True, config_path="unused", identity_file="~/.ssh/test_key"
+        ),
+    )
+    with ExitStack() as stack:
+        stack.enter_context(patch("dropkit.lock.LOCK_FILE", tmp_path / "dropkit.lock"))
+        stack.enter_context(
+            patch(
+                "dropkit.main.load_config_and_api",
+                return_value=(SimpleNamespace(config=config), api),
+            )
+        )
+        find = stack.enter_context(
+            patch("dropkit.main.find_user_droplet", return_value=(None, "testuser"))
+        )
+        ssh = stack.enter_context(patch("dropkit.main.add_ssh_host"))
+        vpn = stack.enter_context(patch("dropkit.main.setup_tailscale", return_value="100.64.0.1"))
+        stack.enter_context(patch("dropkit.main.time.sleep"))
+        yield SimpleNamespace(
+            api=api, snapshot=snapshot, droplet=droplet, find=find, ssh=ssh, vpn=vpn
+        )
+
+
+def test_timeout_then_retry_finishes_same_droplet(wake_env):
+    """A slow restore must retain its resources and resume without creating another VM."""
+    env = wake_env
+    env.api.wait_for_droplet_active.side_effect = DropletWaitTimeoutError("timed out")
+    result = runner.invoke(app, ["wake", "test-droplet", "--timeout", "300", "--no-tailscale"])
+    assert result.exit_code == 1
+    assert "DigitalOcean may still be restoring" in result.output
+    assert "dropkit info test-droplet" in result.output
+    assert "dropkit wake test-droplet --timeout 300 --no-tailscale" in result.output
+    env.ssh.assert_not_called()
+    env.api.delete_snapshot.assert_not_called()
+    env.api.wait_for_droplet_active.assert_called_once_with(123, timeout=300)
+
+    env.find.return_value = (env.droplet, "testuser")
+    env.api.wait_for_droplet_active.side_effect = None
+    result = runner.invoke(app, ["wake", "test-droplet"], input="no\n")
+    assert result.exit_code == 0, result.output
+    assert "Resuming wake" in result.output
+    assert "ssh dropkit.test-droplet" in result.output
+    env.api.create_droplet_from_snapshot.assert_called_once()
+    env.api.wait_for_droplet_active.assert_called_with(123, timeout=900)
+    env.ssh.assert_called_once_with(
+        config_path="unused",
+        host_name="dropkit.test-droplet",
+        hostname="192.0.2.1",
+        user="testuser",
+        identity_file="~/.ssh/test_key",
+    )
+    env.api.delete_snapshot.assert_not_called()
+
+
+@pytest.mark.parametrize("status", ["new", "active"])
+def test_resume_retains_tailscale_setup_and_snapshot_prompt(wake_env, status):
+    env = wake_env
+    env.droplet["status"] = status
+    env.find.return_value = (env.droplet, "testuser")
+    env.snapshot["tags"].append("tailscale-lockdown")
+    result = runner.invoke(app, ["wake", "test-droplet", "--timeout", "1800"], input="yes\n")
+    assert result.exit_code == 0, result.output
+    env.api.create_droplet_from_snapshot.assert_not_called()
+    env.api.wait_for_droplet_active.assert_called_once_with(123, timeout=1800)
+    env.vpn.assert_called_once()
+    env.api.delete_snapshot.assert_called_once_with(456)
+
+
+def test_resume_no_tailscale(wake_env):
+    env = wake_env
+    env.find.return_value = (env.droplet, "testuser")
+    env.snapshot["tags"].append("tailscale-lockdown")
+    result = runner.invoke(app, ["wake", "test-droplet", "--no-tailscale"], input="no\n")
+    assert result.exit_code == 0, result.output
+    env.vpn.assert_not_called()
+
+
+@pytest.mark.parametrize("image", [{"id": 999}, {}, {"id": None}])
+def test_refuse_unrelated_existing_droplet(wake_env, image):
+    env = wake_env
+    env.droplet["image"] = image
+    env.find.return_value = (env.droplet, "testuser")
+    result = runner.invoke(app, ["wake", "test-droplet"])
+    assert result.exit_code == 1
+    assert "was not restored from this snapshot" in " ".join(result.output.split())
+    assert "Destroy or rename" not in result.output
+    env.api.create_droplet_from_snapshot.assert_not_called()
+    env.api.wait_for_droplet_active.assert_not_called()
+    env.api.delete_snapshot.assert_not_called()
+    env.ssh.assert_not_called()
+
+
+@pytest.mark.parametrize("status", ["off", "error", "archive"])
+def test_refuse_resume_when_droplet_cannot_become_ready(wake_env, status):
+    env = wake_env
+    env.droplet["status"] = status
+    env.find.return_value = (env.droplet, "testuser")
+    result = runner.invoke(app, ["wake", "test-droplet"])
+    assert result.exit_code == 1
+    assert "Cannot resume wake" in result.output
+    if status == "off":
+        assert "dropkit on test-droplet" in result.output
+    env.api.wait_for_droplet_active.assert_not_called()
+    env.api.create_droplet_from_snapshot.assert_not_called()
+
+
+def test_missing_snapshot_gives_existing_droplet_recovery_commands(wake_env):
+    env = wake_env
+    env.find.return_value = (env.droplet, "testuser")
+    env.api.get_snapshot_by_name.return_value = None
+    result = runner.invoke(app, ["wake", "test-droplet"])
+    assert result.exit_code == 1
+    assert "dropkit config-ssh test-droplet" in result.output
+    env.api.create_droplet_from_snapshot.assert_not_called()
+    env.ssh.assert_not_called()
+
+
+def test_api_error_during_wait_keeps_resources(wake_env):
+    env = wake_env
+    env.api.wait_for_droplet_active.side_effect = DigitalOceanAPIError("Network error")
+    result = runner.invoke(app, ["wake", "test-droplet"])
+    assert result.exit_code == 1
+    assert "Network error" in result.output
+    assert "rerun wake to resume setup" in result.output
+    env.ssh.assert_not_called()
+    env.api.delete_snapshot.assert_not_called()
+
+
+@pytest.mark.parametrize("failure", ["missing_ip", "ssh_config"])
+def test_setup_failure_keeps_snapshot_for_retry(wake_env, failure):
+    env = wake_env
+    if failure == "missing_ip":
+        env.droplet["networks"]["v4"] = []
+    else:
+        env.ssh.side_effect = OSError("read-only config")
+    result = runner.invoke(app, ["wake", "test-droplet"])
+    assert result.exit_code == 1
+    assert "snapshot kept" in result.output.lower()
+    env.api.delete_snapshot.assert_not_called()
+
+
+@pytest.mark.parametrize("timeout", ["0", "-1"])
+def test_timeout_must_be_positive(wake_env, timeout):
+    result = runner.invoke(app, ["wake", "test-droplet", "--timeout", timeout])
+    assert result.exit_code == 2
+    wake_env.api.create_droplet_from_snapshot.assert_not_called()
+
+
+def test_wait_times_out_at_deadline():
+    api = DigitalOceanAPI("test-token")
+    with (
+        patch.object(api, "get_droplet", return_value={"status": "new"}),
+        patch("time.monotonic", side_effect=[0, 0, 300]),
+        patch("time.sleep"),
+        pytest.raises(DropletWaitTimeoutError, match="waited 300s"),
+    ):
+        api.wait_for_droplet_active(123)
+
+
+def test_wait_allows_restore_longer_than_five_minutes():
+    api = DigitalOceanAPI("test-token")
+    ready = {"status": "active"}
+    with (
+        patch.object(api, "get_droplet", side_effect=[{"status": "new"}, ready]),
+        patch("time.monotonic", side_effect=[0, 395]),
+        patch("time.sleep"),
+    ):
+        assert api.wait_for_droplet_active(123, timeout=900) == ready
+
+
+def test_error_state_is_not_a_timeout():
+    api = DigitalOceanAPI("test-token")
+    with (
+        patch.object(api, "get_droplet", return_value={"status": "error"}),
+        pytest.raises(DigitalOceanAPIError, match="error state") as exc,
+    ):
+        api.wait_for_droplet_active(123)
+    assert not isinstance(exc.value, DropletWaitTimeoutError)

From 65c9436bec02722db6234d7e63f9371b615051ba Mon Sep 17 00:00:00 2001
From: Filipe Casal <fcasal@users.noreply.github.com>
Date: Tue, 8 Sep 2026 18:10:54 +0100
Subject: [PATCH 2/2] Preserve Tailscale SSH configuration when resuming wake

---
 dropkit/main.py    | 32 ++++++++++----------
 tests/test_wake.py | 73 ++++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 86 insertions(+), 19 deletions(-)

diff --git a/dropkit/main.py b/dropkit/main.py
index b34cad5..ad5561d 100644
--- a/dropkit/main.py
+++ b/dropkit/main.py
@@ -4050,7 +4050,7 @@ def wake(
 
     If the original droplet had Tailscale lockdown enabled, this command will
     re-setup Tailscale after the droplet becomes active. Use --no-tailscale to
-    skip this and keep public SSH access.
+    skip this. When resuming, existing Tailscale SSH access is preserved.
 
     Use 'dropkit destroy <name>' to delete a hibernated snapshot without restoring.
     """
@@ -4222,14 +4222,18 @@ def wake(
             try:
                 console.print("[dim]Configuring SSH...[/dim]")
                 ssh_hostname = get_ssh_hostname(droplet_name)
-                add_ssh_host(
-                    config_path=config.ssh.config_path,
-                    host_name=ssh_hostname,
-                    hostname=ip_address,
-                    user=username,
-                    identity_file=config.ssh.identity_file,
-                )
-                console.print("[green]✓[/green] SSH config updated")
+                if existing_droplet and is_droplet_tailscale_locked(config, droplet_name):
+                    # A previous wake may have blocked public SSH before being interrupted.
+                    console.print("[green]✓[/green] Preserved Tailscale SSH config")
+                else:
+                    add_ssh_host(
+                        config_path=config.ssh.config_path,
+                        host_name=ssh_hostname,
+                        hostname=ip_address,
+                        user=username,
+                        identity_file=config.ssh.identity_file,
+                    )
+                    console.print("[green]✓[/green] SSH config updated")
             except Exception as e:
                 console.print(f"[red]Could not update SSH config: {e}[/red]")
                 console.print(
@@ -4244,10 +4248,7 @@ def wake(
             if no_tailscale:
                 console.print()
                 console.print("[yellow]⚠[/yellow] Original droplet had Tailscale lockdown enabled.")
-                console.print(
-                    "[dim]Skipping Tailscale setup (--no-tailscale). "
-                    "Public SSH access available.[/dim]"
-                )
+                console.print("[dim]Skipping Tailscale setup (--no-tailscale).[/dim]")
                 console.print(
                     f"[dim]Enable Tailscale later with: "
                     f"[cyan]dropkit enable-tailscale {droplet_name}[/cyan][/dim]"
@@ -4265,10 +4266,7 @@ def wake(
                 tailscale_ip = setup_tailscale(ssh_hostname, username, config)
 
                 if not tailscale_ip:
-                    console.print(
-                        "[yellow]⚠[/yellow] Tailscale setup incomplete. "
-                        "Public SSH access remains available."
-                    )
+                    console.print("[yellow]⚠[/yellow] Tailscale setup incomplete.")
                     console.print(
                         f"[dim]Complete setup later with: "
                         f"[cyan]dropkit enable-tailscale {droplet_name}[/cyan][/dim]"
diff --git a/tests/test_wake.py b/tests/test_wake.py
index 63cf791..9f41fec 100644
--- a/tests/test_wake.py
+++ b/tests/test_wake.py
@@ -8,7 +8,8 @@
 from typer.testing import CliRunner
 
 from dropkit.api import DigitalOceanAPI, DigitalOceanAPIError, DropletWaitTimeoutError
-from dropkit.main import app
+from dropkit.main import app, setup_tailscale
+from dropkit.ssh_config import add_ssh_host, get_ssh_host_ip
 
 runner = CliRunner()
 
@@ -55,7 +56,13 @@ def wake_env(tmp_path):
         vpn = stack.enter_context(patch("dropkit.main.setup_tailscale", return_value="100.64.0.1"))
         stack.enter_context(patch("dropkit.main.time.sleep"))
         yield SimpleNamespace(
-            api=api, snapshot=snapshot, droplet=droplet, find=find, ssh=ssh, vpn=vpn
+            api=api,
+            snapshot=snapshot,
+            droplet=droplet,
+            find=find,
+            ssh=ssh,
+            vpn=vpn,
+            config=config,
         )
 
 
@@ -113,6 +120,68 @@ def test_resume_no_tailscale(wake_env):
     env.vpn.assert_not_called()
 
 
+@pytest.mark.parametrize("no_tailscale", [False, True])
+def test_retry_after_firewall_lockdown_preserves_tailscale_ssh(wake_env, tmp_path, no_tailscale):
+    env = wake_env
+    env.config.ssh.config_path = str(tmp_path / "ssh_config")
+    env.config.tailscale = SimpleNamespace(lock_down_firewall=True)
+    env.snapshot["tags"].append("tailscale-lockdown")
+    env.ssh.side_effect = add_ssh_host
+    env.vpn.side_effect = setup_tailscale
+    locked = False
+    addresses = []
+
+    def connect(ssh_hostname, verbose):
+        """Simulate an authenticated node that rejects public SSH after lockdown."""
+        address = get_ssh_host_ip(env.config.ssh.config_path, ssh_hostname)
+        addresses.append(address)
+        assert address == ("100.64.0.1" if locked else "192.0.2.1")
+
+    def lockdown(ssh_hostname, verbose):
+        nonlocal locked
+        assert get_ssh_host_ip(env.config.ssh.config_path, ssh_hostname) == "100.64.0.1"
+        locked = True
+        return True
+
+    with (
+        patch("dropkit.main.run_tailscale_up", side_effect=connect),
+        patch("dropkit.main.wait_for_tailscale_ip", return_value="100.64.0.1"),
+        patch("dropkit.main.check_local_tailscale", return_value=True),
+        patch("dropkit.main.lock_down_to_tailscale", side_effect=lockdown),
+        patch("dropkit.main.verify_tailscale_ssh", return_value=True),
+    ):
+        with patch("dropkit.main.Prompt.ask", side_effect=KeyboardInterrupt):
+            result = runner.invoke(app, ["wake", "test-droplet"])
+        assert result.exit_code == 130, result.output
+        assert locked
+        env.api.delete_snapshot.assert_not_called()
+        assert get_ssh_host_ip(env.config.ssh.config_path, "dropkit.test-droplet") == "100.64.0.1"
+
+        env.find.return_value = (env.droplet, "testuser")
+        args = ["wake", "test-droplet"]
+        if no_tailscale:
+            args.append("--no-tailscale")
+        result = runner.invoke(app, args, input="yes\n")
+
+    assert result.exit_code == 0, result.output
+    assert addresses == (["192.0.2.1"] if no_tailscale else ["192.0.2.1", "100.64.0.1"])
+    assert get_ssh_host_ip(env.config.ssh.config_path, "dropkit.test-droplet") == "100.64.0.1"
+    env.api.create_droplet_from_snapshot.assert_called_once()
+    env.api.delete_snapshot.assert_called_once_with(456)
+
+
+def test_fresh_restore_replaces_stale_tailscale_ssh(wake_env, tmp_path):
+    env = wake_env
+    env.config.ssh.config_path = str(tmp_path / "ssh_config")
+    add_ssh_host(env.config.ssh.config_path, "dropkit.test-droplet", "100.64.0.2", "testuser")
+    env.ssh.side_effect = add_ssh_host
+
+    result = runner.invoke(app, ["wake", "test-droplet"], input="no\n")
+
+    assert result.exit_code == 0, result.output
+    assert get_ssh_host_ip(env.config.ssh.config_path, "dropkit.test-droplet") == "192.0.2.1"
+
+
 @pytest.mark.parametrize("image", [{"id": 999}, {}, {"id": None}])
 def test_refuse_unrelated_existing_droplet(wake_env, image):
     env = wake_env