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.


Changed files:
fiftyone/server/main.py, fiftyone/core/session.py, fiftyone/core/client.py

diff --git a/fiftyone/core/client.py b/fiftyone/core/client.py
index 049fed340..e3638d6c2 100644
--- a/fiftyone/core/client.py
+++ b/fiftyone/core/client.py
@@ -114,6 +114,9 @@ class HasClient(object):
                 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 @@ class HasClient(object):
     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
index 65b0ae05c..dae35f06d 100644
--- a/fiftyone/core/session.py
+++ b/fiftyone/core/session.py
@@ -285,6 +285,7 @@ class Session(foc.HasClient):
         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 @@ class Session(foc.HasClient):
         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 @@ class Session(foc.HasClient):
                 )
             )
 
+    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
index 0dd5d54c3..83027e510 100644
--- a/fiftyone/server/main.py
+++ b/fiftyone/server/main.py
@@ -415,6 +415,8 @@ class StateHandler(tornado.websocket.WebSocketHandler):
         """
         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):