Simplify GeneratorExit handling - let it propagate and rely on finally block

app.py

@@ -657,8 +657,6 @@ with gr.Blocks(title="LLM Inference with ZeroGPU") as demo:
             cancel_btn: gr.update(visible=True),
         }
 
-        ui_reset_done = False
-
         try:
             # 2. Call the backend and stream updates
             backend_args = [user_msg, chat_history] + list(args)
@@ -667,16 +665,6 @@ with gr.Blocks(title="LLM Inference with ZeroGPU") as demo:
                     chat: response_chunk[0],
                     dbg: response_chunk[1],
                 }
-        except GeneratorExit:
-            # Handle Gradio's cancellation signal gracefully
-            print("Generation cancelled by user.")
-            # Immediately reset UI state, then re-raise to properly close generator
-            yield {
-                txt: gr.update(interactive=True),
-                submit_btn: gr.update(interactive=True),
-                cancel_btn: gr.update(visible=False),
-            }
-            raise  # Re-raise GeneratorExit to properly close the generator
         except Exception as e:
             print(f"An error occurred during generation: {e}")
             # If an error happens, add it to the chat history to inform the user.
@@ -687,14 +675,12 @@ with gr.Blocks(title="LLM Inference with ZeroGPU") as demo:
             yield {chat: error_history}
         finally:
             # 3. ALWAYS reset the UI to an "idle" state upon completion, error, or cancellation.
-            # Skip if GeneratorExit was raised (already handled above)
-            if not isinstance(sys.exc_info()[1], GeneratorExit):
-                print("Resetting UI state.")
-                yield {
-                    txt: gr.update(interactive=True),
-                    submit_btn: gr.update(interactive=True),
-                    cancel_btn: gr.update(visible=False),
-                }
+            print("Resetting UI state.")
+            yield {
+                txt: gr.update(interactive=True),
+                submit_btn: gr.update(interactive=True),
+                cancel_btn: gr.update(visible=False),
+            }
 
     def set_cancel_flag():
         """Called by the cancel button, sets the global event."""