Bun.serve: close WebSockets with 1001 on server.stop() and let stop(true) follow stop(false)

packages/bun-uws/src/App.h

@@ -385,6 +385,29 @@ struct TemplatedApp {
         return std::move(*this);
     }
 
+    /* Send a close frame to every open WebSocket and fire its close handler.
+     * Unlike close(), this performs the WebSocket closing handshake so peers
+     * observe the given status (e.g. 1001 Going Away) rather than 1006. */
+    TemplatedApp &&endAllWebSockets(int code, std::string_view message = {}) {
+        /* end() fires the close handler synchronously; user JS there can
+         * terminate() a later socket, which rewrites its ->next into the
+         * loop's closed_head and would derail an in-place walk. Snapshot. */
+        std::vector<us_socket_t *> sockets;
+        for (us_socket_group_t *g : webSocketGroups) {
+            for (struct us_socket_t *s = g->head_sockets; s; s = s->next) {
+                sockets.push_back(s);
+            }
+        }
+        for (us_socket_t *s : sockets) {
+            if (!us_socket_is_closed(s)) {
+                /* USERDATA is erased in the handler slots; see the TopicTree
+                 * cast above. end() no-ops on isShuttingDown. */
+                ((WebSocket<SSL, true, int> *) s)->end(code, message);
+            }
+        }
+        return std::move(*this);
+    }
+
     /** Closes all connections connected to this server which are not sending a request or waiting for a response. Does not close the listen socket. */
     TemplatedApp &&closeIdle() {
         auto *group = httpContext->getSocketGroup();

src/runtime/server/mod.rs

@@ -1632,18 +1632,16 @@ impl<const SSL: bool, const DEBUG: bool> NewServer<SSL, DEBUG> {
             if Self::HAS_H3 && self.h3_app.is_some() {
                 self.unref();
                 self.notify_inspector_server_stopped();
-                if abrupt {
-                    self.flags.insert(ServerFlags::TERMINATED);
-                }
+            }
+            // A prior graceful stop already took the listener. An abrupt stop
+            // still needs to tear down in-flight connections so a "graceful
+            // then force" shutdown can complete.
+            if abrupt {
+                self.end_all_websockets_going_away();
+                self.terminate_app();
             }
             return;
         };
-        if abrupt || (self.pending_requests == 0 && !self.has_active_web_sockets()) {
-            self.unref();
-        }
-        // A graceful stop with work in flight keeps the ref (deinit_if_we_can
-        // unrefs when the drain completes): on Windows uv_run skips I/O with
-        // zero ref'd handles, so unrefing here wedged server.close() teardown.
 
         if !SSL {
             // SAFETY: `listener` is a live uws ListenSocket FFI handle just taken
@@ -1664,31 +1662,72 @@ impl<const SSL: bool, const DEBUG: bool> NewServer<SSL, DEBUG> {
             }
         }
 
+        // Send 1001 Going Away to every open WebSocket so peers see a clean
+        // protocol close instead of a dropped connection (1006), and so the
+        // graceful-stop promise can observe the count reaching zero.
+        self.end_all_websockets_going_away();
+
+        if abrupt || (self.pending_requests == 0 && !self.has_active_web_sockets()) {
+            self.unref();
+        }
+        // A graceful stop with work in flight keeps the ref (deinit_if_we_can
+        // unrefs when the drain completes): on Windows uv_run skips I/O with
+        // zero ref'd handles, so unrefing here wedged server.close() teardown.
+
         if !abrupt {
             // S012: `app::ListenSocket<SSL>` is a ZST opaque — safe deref.
             bun_opaque::opaque_deref_mut(listener).close();
-        } else if !self.flags.contains(ServerFlags::TERMINATED) {
-            if let Some(ws) = self.config.websocket.as_mut() {
-                ws.handler.app = None;
-            }
-            self.flags.insert(ServerFlags::TERMINATED);
-            // `app.close()` synchronously drains every open websocket; their
-            // `on_close` defers call `on_websocket_closed`, which would
-            // dispatch `deinit_if_we_can` through a fresh `&mut NewServer`
-            // while this frame still holds `&mut self`. Hold the re-entrance
-            // guard across the drain so that nested call early-returns —
-            // `stop()` runs `deinit_if_we_can` itself right after this returns.
-            self.deinit_running.set(true);
+        } else {
+            self.terminate_app();
+        }
+    }
+
+    fn end_all_websockets_going_away(&mut self) {
+        if !self.has_active_web_sockets() {
+            return;
+        }
+        // node:http `Server#close()` must leave upgraded sockets to the user
+        // (Node only stops accepting and closes idle keep-alives); the `ws`
+        // shim tracks its own `clients` set for the user to drain.
+        if !self.config.on_node_http_request.is_empty() {
+            return;
+        }
+        let Some(app) = self.app else { return };
+        // `end()` fires the close handler synchronously; its `on_close` defer
+        // calls `on_websocket_closed`, which would dispatch `deinit_if_we_can`
+        // through a fresh `&mut NewServer` while this frame still holds
+        // `&mut self`. Hold the re-entrance guard across the drain so the
+        // nested call early-returns; `stop()` runs the idle pass afterwards.
+        // Save/restore so a nested `server.stop(true)` from a close handler
+        // cannot clear the outer frame's guard.
+        let prev = self.deinit_running.replace(true);
+        // S012: `NewApp<SSL>` is a ZST opaque — safe `*mut → &mut` deref.
+        bun_opaque::opaque_deref_mut(app).end_all_websockets(1001, b"Server closed");
+        self.deinit_running.set(prev);
+    }
+
+    /// Force-close every connection on the uws app and mark the server
+    /// terminated. Guarded by `TERMINATED` so repeated abrupt stops are no-ops.
+    fn terminate_app(&mut self) {
+        if self.flags.contains(ServerFlags::TERMINATED) {
+            return;
+        }
+        if let Some(ws) = self.config.websocket.as_mut() {
+            ws.handler.app = None;
+        }
+        self.flags.insert(ServerFlags::TERMINATED);
+        if let Some(app) = self.app {
+            let prev = self.deinit_running.replace(true);
             // S012: `NewApp<SSL>` is a ZST opaque — safe `*mut → &mut` deref.
-            bun_opaque::opaque_deref_mut(self.app.unwrap()).close();
-            self.deinit_running.set(false);
-            // Only clear after the drain — `on_close` defers reach
-            // `on_websocket_closed` through `handler.server`, so wiping it
-            // earlier would strand the live-socket count and the idle pass
-            // would never see it drained.
-            if let Some(ws) = self.config.websocket.as_mut() {
-                ws.handler.server = None;
-            }
+            bun_opaque::opaque_deref_mut(app).close();
+            self.deinit_running.set(prev);
+        }
+        // Only clear after the drain — `on_close` defers reach
+        // `on_websocket_closed` through `handler.server`, so wiping it
+        // earlier would strand the live-socket count and the idle pass
+        // would never see it drained.
+        if let Some(ws) = self.config.websocket.as_mut() {
+            ws.handler.server = None;
         }
     }
 

src/runtime/server/server_body.rs

@@ -2549,23 +2549,32 @@ where
     pub fn stop_from_js(&mut self, abruptly: Option<JSValue>) -> JSValue {
         let rc = self.get_all_closed_promise(&self.global());
 
-        if self.has_listener() {
-            let abrupt = 'brk: {
-                if let Some(val) = abruptly {
-                    if val.is_boolean() && val.to_boolean() {
-                        break 'brk true;
-                    }
+        let abrupt = 'brk: {
+            if let Some(val) = abruptly {
+                if val.is_boolean() && val.to_boolean() {
+                    break 'brk true;
                 }
-                false
-            };
+            }
+            false
+        };
+        // `!deinit_running`: a `server.stop()` from inside a websocket close
+        // handler fired by an outer `stop()`'s drain would re-enter
+        // `stop_listening` with a fresh `&mut self` under the outer borrow.
+        if self.has_listener()
+            || (abrupt
+                && !self.flags.contains(ServerFlags::TERMINATED)
+                && !self.deinit_running.get())
+        {
             self.stop(abrupt);
         }
 
         rc
     }
 
     pub fn dispose_from_js(&mut self) -> JSValue {
-        if self.has_listener() {
+        if self.has_listener()
+            || (!self.flags.contains(ServerFlags::TERMINATED) && !self.deinit_running.get())
+        {
             self.stop(true);
         }
         JSValue::UNDEFINED
@@ -2742,7 +2751,7 @@ where
     }
 
     pub fn get_all_closed_promise(&mut self, global: &JSGlobalObject) -> JSValue {
-        if !self.has_listener() && self.pending_requests == 0 {
+        if !self.has_listener() && self.pending_requests == 0 && !self.has_active_web_sockets() {
             return JSPromise::resolved_promise(global, JSValue::UNDEFINED).to_js();
         }
         if self.all_closed_promise.has_value() {

src/uws_sys/App.rs

@@ -105,6 +105,20 @@ impl<const SSL: bool> App<SSL> {
         c::uws_app_close_idle(Self::SSL_FLAG, self.as_raw())
     }
 
+    pub fn end_all_websockets(&mut self, code: i32, message: &[u8]) {
+        // SAFETY: `message` is only read for the synchronous duration of the
+        // call; uWS copies it into the close frame before returning.
+        unsafe {
+            c::uws_app_end_all_websockets(
+                Self::SSL_FLAG,
+                self.as_raw(),
+                code,
+                message.as_ptr(),
+                message.len(),
+            )
+        }
+    }
+
     pub fn create(opts: &BunSocketContextOptions) -> Option<*mut Self> {
         // SAFETY: FFI call; uws_create_app returns null on failure.
         let app = unsafe { c::uws_create_app(Self::SSL_FLAG, *opts) };
@@ -513,6 +527,13 @@ pub mod c {
     unsafe extern "C" {
         pub(crate) safe fn uws_app_close(ssl: i32, app: &mut uws_app_s);
         pub(crate) safe fn uws_app_close_idle(ssl: i32, app: &mut uws_app_s);
+        pub(crate) fn uws_app_end_all_websockets(
+            ssl: i32,
+            app: &mut uws_app_s,
+            code: c_int,
+            message: *const u8,
+            length: usize,
+        );
         // safe: `&mut uws_app_s` is ABI-identical to a non-null `*mut`;
         // `handler`/`user_data` are stored opaquely (never dereferenced by the
         // C++ shim itself) — no preconditions on this call.

src/uws_sys/libuwsockets.cpp

@@ -401,6 +401,21 @@ extern "C"
     }
   }
 
+  void uws_app_end_all_websockets(int ssl, uws_app_t *app, int code, const char *message, size_t length)
+  {
+    std::string_view msg = message ? std::string_view(message, length) : std::string_view();
+    if (ssl)
+    {
+      uWS::SSLApp *uwsApp = (uWS::SSLApp *)app;
+      uwsApp->endAllWebSockets(code, msg);
+    }
+    else
+    {
+      uWS::App *uwsApp = (uWS::App *)app;
+      uwsApp->endAllWebSockets(code, msg);
+    }
+  }
+
   void uws_app_set_on_clienterror(int ssl, uws_app_t *app, void (*handler)(void *user_data, int is_ssl, struct us_socket_t *rawSocket, uint8_t errorCode, char *rawPacket, int rawPacketLength), void *user_data)
   {
     if (ssl)