Register the OAuth gate inside the Blocks context

app.py

@@ -917,8 +917,19 @@ def main() -> None:
         return result
 
     def controlled_blocks_exit(blocks, exc_type, exc_value, traceback):
+        nonlocal auth_load_configured
         if blocks is gated_blocks and app_shell is not None:
             app_shell.__exit__(exc_type, exc_value, traceback)
+            if auth_gate is not None and not auth_load_configured:
+                blocks.load(
+                    auth_visibility,
+                    inputs=None,
+                    outputs=[auth_gate, app_shell],
+                    queue=False,
+                    api_name=False,
+                )
+                auth_load_configured = True
+                _trace("phase=frontend.auth_gate event=REGISTERED")
         return original_blocks_exit(blocks, exc_type, exc_value, traceback)
 
     gr.Blocks.__enter__ = controlled_blocks_enter
@@ -965,22 +976,6 @@ def main() -> None:
 
         @wraps(original_launch)
         def controlled_launch(blocks, *args, **kwargs):
-            nonlocal auth_load_configured
-            if (
-                blocks is gated_blocks
-                and auth_gate is not None
-                and app_shell is not None
-                and not auth_load_configured
-            ):
-                blocks.load(
-                    auth_visibility,
-                    inputs=None,
-                    outputs=[auth_gate, app_shell],
-                    queue=False,
-                    api_name=False,
-                )
-                auth_load_configured = True
-                _trace("phase=frontend.auth_gate event=REGISTERED")
             if not queue_configured:
                 controlled_queue(blocks)
             kwargs["ssr_mode"] = False