diff --git a/pyproject.toml b/pyproject.toml
index c64cdb8..d7aa674 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -14,7 +14,7 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
-clawhub = "clawhub>=0.1.0"
+clawhub = ["clawhub>=0.1.0"]
 
 [project.scripts]
 ocbs = "ocbs.cli:main"
@@ -32,4 +32,4 @@ dev = [
 
 [tool.pytest.ini_options]
 testpaths = ["tests"]
-python_files = ["test_*.py"]
+python_files = ["test_*.py"]
\ No newline at end of file
diff --git a/src/ocbs/serve.py b/src/ocbs/serve.py
index 7fec8fe..57a8483 100644
--- a/src/ocbs/serve.py
+++ b/src/ocbs/serve.py
@@ -362,6 +362,53 @@ def _get_html_page(self, token: str, checkpoint_info: dict, expires_at: datetime
         step1_button_text = "Step 1: I received this - start changes"
         step1_button_disabled = ""
         work_underway_display = "none"
+
+        # Get raw scope (before HTML escaping) to determine what data exists
+        raw_scope = checkpoint_info.get('scope', '')
+
+        # Stage 2: Partial restore (config, or config+session if workspace exists)
+        if raw_scope == "config":
+            stage2_label = "Restore Configuration & Restart"
+        elif raw_scope == "config+session":
+            stage2_label = "Restore Config + Sessions & Restart"
+        elif raw_scope == "config+session+workspace":
+            stage2_label = "Restore Config + Sessions & Restart"
+        else:
+            stage2_label = "Restore Backup & Restart"
+
+        # Stage 3: Full restore (only for backups with workspace/session data)
+        has_extended_data = "session" in raw_scope or "workspace" in raw_scope
+        stage3_label = "Full Restore & Restart"
+
+        # Build stage2 button HTML
+        stage2_html = f"""
+            <form action="/restore" method="POST">
+                <input type="hidden" name="token" value="{token}">
+                <input type="hidden" name="restore_scope" value="partial">
+                <button type="submit" class="btn btn-danger" {step1_button_disabled}
+                        onclick="return confirm('Are you sure you want to RESTORE? This will revert configuration changes!');">
+                    <span class="step-label step2">Step 2</span>
+                    &#9888;&#65039; {stage2_label}
+                    <span class="spinner" id="restore-spinner"></span>
+                </button>
+            </form>
+        """
+
+        # Build stage3 button HTML (only if session/workspace data exists)
+        stage3_html = ""
+        if has_extended_data:
+            stage3_html = f"""
+            <form action="/restore" method="POST">
+                <input type="hidden" name="token" value="{token}">
+                <input type="hidden" name="restore_scope" value="full">
+                <button type="submit" class="btn btn-danger" {step1_button_disabled}
+                        onclick="return confirm('Are you sure you want to do a FULL RESTORE? This will revert ALL changes including workspace data!');">
+                    <span class="step-label step2">Step 3</span>
+                    &#9888;&#65039; {stage3_label}
+                    <span class="spinner" id="full-restore-spinner"></span>
+                </button>
+            </form>
+        """
         
         if is_restored:
             status_message = """
@@ -615,19 +662,13 @@ def _get_html_page(self, token: str, checkpoint_info: dict, expires_at: datetime
                 <input type="hidden" name="token" value="{token}">
                 <button type="submit" class="btn btn-restart" {step1_button_disabled}
                         onclick="return confirm('Restart the gateway without restoring? This may help resolve minor issues.');">
-                    🔄 Restart Gateway
-                </button>
-            </form>
-            
-            <form action="/restore" method="POST">
-                <input type="hidden" name="token" value="{token}">
-                <button type="submit" class="btn btn-danger" {step1_button_disabled} 
-                        onclick="return confirm('Are you sure you want to RESTORE? This will revert all changes!');">
-                    <span class="step-label step2">Step 2</span>
-                    🔴 Restore Backup & Restart
-                    <span class="spinner" id="restore-spinner"></span>
+                    &#8635; Restart Gateway Only
                 </button>
             </form>
+
+            {stage2_html}
+
+            {stage3_html}
         </div>
         
         <p class="expiry">Link expires in {hours}h {minutes}m</p>
@@ -693,13 +734,16 @@ def _get_html_page(self, token: str, checkpoint_info: dict, expires_at: datetime
             }});
         }});
         
-        // Handle restore form
-        document.querySelector('form[action="/restore"]').addEventListener('submit', function(e) {{
-            if (!confirm('Are you sure? This will restore from the checkpoint.')) {{
-                e.preventDefault();
-            }} else {{
-                document.getElementById('restore-spinner').style.display = 'inline-block';
-            }}
+        // Handle restore forms (both partial and full)
+        document.querySelectorAll('form[action="/restore"]').forEach(function(form) {{
+            form.addEventListener('submit', function(e) {{
+                const scope = this.querySelector('input[name="restore_scope"]').value;
+                const spinnerId = (scope === 'full') ? 'full-restore-spinner' : 'restore-spinner';
+                const spinner = document.getElementById(spinnerId);
+                if (spinner) {{
+                    spinner.style.display = 'inline-block';
+                }}
+            }});
         }});
     </script>
 </body>