Repository: voxel51/fiftyone

Issue:
[BUG] session.wait() does not detect closing the web interface
### System information
-   **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: Manjaro Linux
-   **FiftyOne installed from (pip or source)**: pip
-   **FiftyOne version (run `fiftyone --version`)**: v0.8.0.1,
-   **Python version**: 3.8.8

### Describe the problem
If I launch the FiftyOne app from a non-interactive Python script and call `session.wait()`, the web app is automatically opened in my (already open) browser, but closing it doesn't seem to trigger anything. The script blocks forever until I shut it down manually with CTRL + C, and then the code below `session.wait()` is never executed.

From the docs this seems like a pretty standard use-case, so I suspect I may be doing something wrong? Any help would be greatly appreciated! I intend to use the app as a somewhat standalone interface; I would strongly prefer not having to wrap it in a notebook. Thanks!

### Code to reproduce issue

```python
    directory = "some_directory"
    dataset = fo.Dataset.from_dir(directory, fo.types.ImageDirectory, name="Local data")
    session = fo.launch_app(dataset)
    session.wait()
    dataset.save()
    print("Saved FiftyOne dataset")
```

### What areas of FiftyOne does this bug affect?

-   [x] `App`: FiftyOne application issue
-   [ ] `Core`: Core `fiftyone` Python library issue
-   [ ] `Server`: Fiftyone server issue

### Willingness to contribute

The FiftyOne Community encourages bug fix contributions. Would you or another
member of your organization be willing to contribute a fix for this bug to the
FiftyOne codebase?

-   [ ] Yes. I can contribute a fix for this bug independently.
-   [ ] Yes. I would be willing to contribute a fix for this bug with guidance
        from the FiftyOne community.
-   [x] No. I cannot contribute a bug fix at this time.



Discussion:
Ah, I've now tried launching it with `desktop=True` and that works as intended. Is this behaviour not supported with the web interface? It would make complete sense if not, but then perhaps it would be nice to add that to the docs :)
Hi @jneeven. Thanks for the issue! Browser support for `session.wait()` was previously overlooked, apologies. I'll be adding support, barring any browser limitations. Stay tuned!

Resolving patch:
diff --git a/fiftyone/core/client.py b/fiftyone/core/client.py
--- a/fiftyone/core/client.py
+++ b/fiftyone/core/client.py
@@ -114,6 +114,9 @@ async def connect():
                 if event == "reload":
                     self.on_reload()
 
+                if event == "close":
+                    self.on_close()
+
         def run_client():
             io_loop = IOLoop(make_current=True)
             io_loop.run_sync(connect)
@@ -169,6 +172,12 @@ def on_capture(self, data):
     def _capture(self, data):
         raise NotImplementedError("subclasses must implement _capture()")
 
+    def on_close(self):
+        self._close()
+
+    def _close(self):
+        raise NotImplementedError("subclasses must implement _close()")
+
     def on_reactivate(self, data):
         self._reactivate(data)
 
diff --git a/fiftyone/core/session.py b/fiftyone/core/session.py
--- a/fiftyone/core/session.py
+++ b/fiftyone/core/session.py
@@ -285,6 +285,7 @@ def __init__(
         self._plots = None
         self._port = port
         self._remote = remote
+        self._wait_closed = False
 
         # Maintain a reference to prevent garbage collection
         self._get_time = time.perf_counter
@@ -845,20 +846,19 @@ def no_show(self):
         return fou.SetAttributes(self, _auto=False)
 
     def wait(self):
-        """Blocks execution until the session is closed by the user.
+        """Blocks execution until the App is closed by the user.
 
-        For local sessions, this will wait until the App is closed by the user.
-
-        For remote sessions, this will wait until the server shuts down, which
-        typically requires interrupting the calling process with Ctrl-C.
+        Closing a desktop App will always exit the wait. For browser Apps, all
+        connected windows (tabs) must be closed.
         """
+        if self._context != focx._NONE:
+            raise RuntimeError("Notebook sessions cannot wait")
+
         try:
             if self._remote or not self._desktop:
-                try:
-                    _server_services[self._port].wait()
-                except:
-                    while True:
-                        time.sleep(1)
+                self._wait_closed = False
+                while not self._wait_closed:
+                    time.sleep(1)
             else:
                 self._app_service.wait()
         except KeyboardInterrupt:
@@ -913,6 +913,9 @@ def _capture(self, data):
                 )
             )
 
+    def _close(self):
+        self._wait_closed = True
+
     def _base_url(self):
         if self._context == focx._COLAB:
             # pylint: disable=no-name-in-module,import-error
diff --git a/fiftyone/server/main.py b/fiftyone/server/main.py
--- a/fiftyone/server/main.py
+++ b/fiftyone/server/main.py
@@ -415,6 +415,8 @@ def on_close(self):
         """
         StateHandler.clients.remove(self)
         StateHandler.app_clients.discard(self)
+        if not StateHandler.app_clients:
+            _write_message({"type": "close"}, session=True)
 
     @_catch_errors
     async def on_message(self, message):


Test patch: