3 files changed, 13 insertions, 11 deletions
diff --git a/Userland/Services/WebDriver/Session.cpp b/Userland/Services/WebDriver/Session.cpp
index f18303eaae..011d37d454 100644
--- a/Userland/Services/WebDriver/Session.cpp
+++ b/Userland/Services/WebDriver/Session.cpp
@@ -61,7 +61,7 @@ ErrorOr<NonnullRefPtr<Core::LocalServer>> Session::create_server(NonnullRefPtr<S
     server->listen(*m_web_content_socket_path);
 
     server->on_accept = [this, promise](auto client_socket) {
-        auto maybe_connection = adopt_nonnull_ref_or_enomem(new (nothrow) WebContentConnection(move(client_socket), m_client, session_id()));
+        auto maybe_connection = adopt_nonnull_ref_or_enomem(new (nothrow) WebContentConnection(move(client_socket)));
         if (maybe_connection.is_error()) {
             // Use of MUST in this function is safe, as our promise callback can never error out.
             MUST(promise->resolve(maybe_connection.release_error()));
@@ -72,6 +72,12 @@ ErrorOr<NonnullRefPtr<Core::LocalServer>> Session::create_server(NonnullRefPtr<S
         auto web_content_connection = maybe_connection.release_value();
 
         auto window_handle = web_content_connection->get_window_handle();
+        web_content_connection->on_close = [this, window_handle]() {
+            dbgln_if(WEBDRIVER_DEBUG, "Window {} was closed remotely.", window_handle);
+            m_windows.remove(window_handle);
+            if (m_windows.is_empty())
+                m_client->close_session(session_id());
+        };
         m_windows.set(window_handle, Session::Window { window_handle, move(web_content_connection) });
 
         if (m_current_window_handle.is_empty())
diff --git a/Userland/Services/WebDriver/WebContentConnection.cpp b/Userland/Services/WebDriver/WebContentConnection.cpp
index f7ae99edb9..8b6aeb2f85 100644
--- a/Userland/Services/WebDriver/WebContentConnection.cpp
+++ b/Userland/Services/WebDriver/WebContentConnection.cpp
@@ -9,17 +9,15 @@
 
 namespace WebDriver {
 
-WebContentConnection::WebContentConnection(NonnullOwnPtr<Core::LocalSocket> socket, NonnullRefPtr<Client> client, unsigned session_id)
+WebContentConnection::WebContentConnection(NonnullOwnPtr<Core::LocalSocket> socket)
     : IPC::ConnectionFromClient<WebDriverClientEndpoint, WebDriverServerEndpoint>(*this, move(socket), 1)
-    , m_client(move(client))
-    , m_session_id(session_id)
 {
 }
 
 void WebContentConnection::die()
 {
-    dbgln_if(WEBDRIVER_DEBUG, "Session {} was closed remotely. Shutting down...", m_session_id);
-    m_client->close_session(m_session_id);
+    if (on_close)
+        on_close();
 }
 
 }
diff --git a/Userland/Services/WebDriver/WebContentConnection.h b/Userland/Services/WebDriver/WebContentConnection.h
index a04429631b..6c31a45344 100644
--- a/Userland/Services/WebDriver/WebContentConnection.h
+++ b/Userland/Services/WebDriver/WebContentConnection.h
@@ -18,13 +18,11 @@ class WebContentConnection
     : public IPC::ConnectionFromClient<WebDriverClientEndpoint, WebDriverServerEndpoint> {
     C_OBJECT_ABSTRACT(WebContentConnection)
 public:
-    WebContentConnection(NonnullOwnPtr<Core::LocalSocket> socket, NonnullRefPtr<Client> client, unsigned session_id);
+    WebContentConnection(NonnullOwnPtr<Core::LocalSocket> socket);
 
-    virtual void die() override;
+    Function<void()> on_close;
 
-private:
-    NonnullRefPtr<Client> m_client;
-    unsigned m_session_id { 0 };
+    virtual void die() override;
 };
 
 }