Update app.py

app.py

@@ -12,6 +12,7 @@ from pathlib import Path
 from PIL import Image
 
 # --- Safe Input Mocking ---
+# Fix: Mock input to prevent hanging on server consoles
 builtins.input = lambda *args: "y"
 
 # GenAI & ADK Imports
@@ -28,7 +29,7 @@ except ImportError as e:
     Sam3Model = None
     Sam3Processor = None
     root_agent = None
-    AnalysisDeps = None  # <--- FIX: Indented inside except block 
+    AnalysisDeps = None  # FIX: Correctly indented inside except block
 
 # Optional: Distinctipy for better colors
 try:
@@ -268,30 +269,41 @@ async def run_analysis(image_path_str, user_prompt, session_id_state):
         yield [{"role": "assistant", "content": error_msg}], None, None, [], None, waiting_df, waiting_df, waiting_df, *empty_slider_updates
         return
 
-    # FIX: Add safety check for AnalysisDeps
+    # FIX: Safety check for AnalysisDeps availability
     if AnalysisDeps is None:
         error_msg = "❌ Project imports failed. 'AnalysisDeps' is missing. Check your 'cellemetry' package installation."
         yield [{"role": "assistant", "content": error_msg}], None, None, [], None, waiting_df, waiting_df, waiting_df, *empty_slider_updates
         return
 
-    deps = AnalysisDeps(
-        sam_model=MODEL_CACHE["model"],
-        sam_processor=MODEL_CACHE["processor"],
-        image_path=image_path,
-        device=MODEL_CACHE["device"],
-        pixel_size_microns=None 
-    )
-    
-    global ACTIVE_RUNNER
-    ACTIVE_RUNNER = InMemoryRunner(agent=root_agent, app_name="cellemetry_demo")
-    
-    session = await ACTIVE_RUNNER.session_service.create_session(
-        app_name="cellemetry_demo",
-        user_id="demo_user",
-        state=deps.to_state_dict()
-    )
-    
-    session_id = session.id
+    # FIX: Wrap Agent/Runner initialization in try/except to catch crashes here
+    try:
+        deps = AnalysisDeps(
+            sam_model=MODEL_CACHE["model"],
+            sam_processor=MODEL_CACHE["processor"],
+            image_path=image_path,
+            device=MODEL_CACHE["device"],
+            pixel_size_microns=None 
+        )
+        
+        global ACTIVE_RUNNER
+        if root_agent is None:
+            raise ValueError("Root agent is not loaded.")
+            
+        ACTIVE_RUNNER = InMemoryRunner(agent=root_agent, app_name="cellemetry_demo")
+        
+        session = await ACTIVE_RUNNER.session_service.create_session(
+            app_name="cellemetry_demo",
+            user_id="demo_user",
+            state=deps.to_state_dict()
+        )
+        session_id = session.id
+        
+    except Exception as e:
+        error_msg = f"❌ Agent Initialization Failed: {str(e)}"
+        print(error_msg)
+        yield [{"role": "assistant", "content": error_msg}], None, None, [], None, waiting_df, waiting_df, waiting_df, *empty_slider_updates
+        return
+
     image_bytes = image_path.read_bytes()
     content = types.Content(
         role="user",
@@ -394,16 +406,22 @@ async def unified_chat_handler(message, history, session_id, current_img_path):
         yield history, session_id, image_path, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), *empty_slider_updates, None, gr.update(visible=False), gr.update(visible=True), gr.update(visible=False)
         
         final_result = None
-        async for result in run_analysis(image_path, user_text, session_id):
-            final_result = result
-            updated_history = result[0].copy()
-            if files and len(updated_history) > 0:
-                # Ensure the user message in history also uses the encoded path
-                updated_history[0] = {"role": "user", "content": f"![](file={display_path})\n\n{user_text}"}
-            
-            # Yield Loop: Pass gr.update() to prevent flickering
-            yield (updated_history, result[1], image_path, *result[2:], None, gr.update(), gr.update(), gr.update())
-        
+        # FIX: Catch potential generator crashes here to prevent asyncio loop death
+        try:
+            async for result in run_analysis(image_path, user_text, session_id):
+                final_result = result
+                updated_history = result[0].copy()
+                if files and len(updated_history) > 0:
+                    # Ensure the user message in history also uses the encoded path
+                    updated_history[0] = {"role": "user", "content": f"![](file={display_path})\n\n{user_text}"}
+                
+                # Yield Loop: Pass gr.update() to prevent flickering
+                yield (updated_history, result[1], image_path, *result[2:], None, gr.update(), gr.update(), gr.update())
+        except Exception as e:
+            history.append({"role": "assistant", "content": f"❌ Critical Error: {str(e)}"})
+            yield history, session_id, image_path, None, gr.CheckboxGroup(), None, waiting_df, waiting_df, waiting_df, *empty_slider_updates, None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)
+            return
+
         if final_result:
             updated_history = final_result[0].copy()
             if files and len(updated_history) > 0:
@@ -420,7 +438,7 @@ async def unified_chat_handler(message, history, session_id, current_img_path):
         yield history, session_id, current_img_path, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), *empty_slider_updates, None, gr.update(), gr.update(), gr.update()
         
         if not ACTIVE_RUNNER:
-            history[-1]["content"] = "⚠️ Session expired."
+            history[-1]["content"] = "⚠️ Session expired or Agent not initialized."
             yield history, None, current_img_path, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), *empty_slider_updates, None, gr.update(), gr.update(), gr.update()
             return