Update app.py

app.py

@@ -36,30 +36,56 @@ def get_current_status():
         return f"❌ Connection Error: {str(e)}"
 
 def manage_instance(action, confirmed=False):
-    # ... (Keep existing setup code) ...
+    """Handles Activate and Shelve actions with milestone updates."""
+    clean_link = "chat.mch250095.projects.jetstream-cloud.org"
+    full_url = f"https://{clean_link}/"
+    instance_id = os.getenv("OS_INSTANCE_ID")
+    
     try:
-        # ... (Keep connection and server find logic) ...
+        # CRITICAL FIX: Define 'conn' and 'server' at the very beginning of the try block
+        conn = get_conn()
+        server = conn.compute.find_server(instance_id)
+        
+        if not server:
+            yield ("❌ Error: Instance ID not found.", gr.update(visible=False))
+            return
 
         if action == "activate":
-            # Milestone 1...
-            # [Previous logic for unshelving]
+            yield ("⏳ Milestone 1/3: Un-shelving Begin...", gr.update(visible=False))
+            if server.status == "SHELVED_OFFLOADED":
+                conn.compute.unshelve_server(server)
             
-            # Milestone 2: Unshelved with extended 10-minute (600s) timeout
+            # Milestone 2: Extended 10-minute (600s) timeout
             try:
                 conn.compute.wait_for_server(server, status='ACTIVE', wait=600)
                 yield ("✅ Milestone 2/3: Unshelved! VM is now ACTIVE.", gr.update(visible=False))
             except exceptions.ResourceTimeout:
                 yield ("⚠️ Milestone 2/3: Booting is taking a long time (10+ mins).", gr.update(visible=False))
+                return
 
-            # Milestone 3...
-            # [Previous polling logic]
+            # Milestone 3: Web service check
+            yield ("🔄 Milestone 3/3: Waiting for software...", gr.update(visible=False))
+            for i in range(24): 
+                try:
+                    response = requests.get(full_url, timeout=3)
+                    if response.status_code < 400:
+                        yield ("🎉 Success! Ready.", gr.update(value=f"### [Visit: {clean_link}]({full_url})", visible=True))
+                        return
+                except:
+                    pass
+                yield (f"🔄 Milestone 3/3: Software booting (Attempt {i+1}/24)...", gr.update(visible=False))
+                time.sleep(5)
+            
+            yield ("⚠️ Service slow.", gr.update(value=f"### [Manual Link: {clean_link}]({full_url})", visible=True))
 
         elif action == "shelve":
-            # [Confirmation check]
-            yield ("💤 Shelving Begin...", gr.update(visible=False))
+            if not confirmed:
+                yield ("🛑 Shelve aborted: Confirmation required.", gr.update(visible=False))
+                return
+
+            yield ("💤 Shelving Begin... (Saving and Stopping)", gr.update(visible=False))
             conn.compute.shelve_server(server)
-            
-            # Also extend the shelve timeout to 10 minutes to be safe
+            # Extended timeout for shelving as well
             conn.compute.wait_for_server(server, status='SHELVED_OFFLOADED', wait=600)
             yield ("🛌 Instance is now SHELVED.", gr.update(visible=False))
 
@@ -82,10 +108,10 @@ with gr.Blocks() as demo:
     status_label = gr.Label(value="Ready")
     link_box = gr.Markdown(visible=False)
 
-    # UI Logic: Enable button only if checkbox is ticked
+    # UI Logic
     confirm_check.change(fn=lambda x: gr.update(interactive=x), inputs=confirm_check, outputs=shelve_btn)
 
-    # ACTION CLICKS: Using gr.State to avoid Lambda Generator issues
+    # ACTION CLICKS: Direct function reference with gr.State
     activate_btn.click(
         fn=manage_instance, 
         inputs=[gr.State("activate")], 
@@ -98,10 +124,10 @@ with gr.Blocks() as demo:
         outputs=[status_label, link_box]
     )
 
-    # Auto-refreshing status
+    # Auto-refreshing status logic for Gradio 5/6
     timer = gr.Timer(30)
     timer.tick(fn=get_current_status, outputs=live_status)
     demo.load(fn=get_current_status, outputs=live_status)
 
 if __name__ == "__main__":
-    demo.launch(theme=gr.themes.Soft())
+    demo.launch(theme=gr.themes.Soft())