Fix use-after-free in ServerWebSocket.onClose

Move active_connections decrement before ctx.deinit() to avoid use-after-free.
The handler pointer may point inside the WebSocketServerContext, so accessing
handler.active_connections after ctx.deinit() causes AddressSanitizer errors.

Fixes AddressSanitizer: use-after-poison error at line 313.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

packages/bun-types/serve.d.ts

@@ -535,18 +535,41 @@ declare module "bun" {
 
     type BaseRouteValue = Response | false | HTMLBundle | BunFile;
 
+    /**
+     * Route configuration with optional WebSocket handler.
+     * When `websocket` is specified, an `upgrade` handler must also be provided.
+     * @template WebSocketData - Type of data attached to WebSocket connections
+     * @template Path - Route path for typed route parameters (e.g., "/user/:id")
+     * @template HTTPResponse - HTTP handler return type (Response or Response | undefined | void)
+     */
+    type RouteWithWebSocket<WebSocketData, Path extends string = string, HTTPResponse = Response> =
+      | (Partial<Record<HTTPMethod, Handler<BunRequest<Path>, Server<WebSocketData>, HTTPResponse>>> & {
+          websocket: WebSocketHandler<WebSocketData>;
+          /**
+           * Upgrade handler for WebSocket connections.
+           * Required when `websocket` is specified.
+           */
+          upgrade: Handler<BunRequest<Path>, Server<WebSocketData>, Response | undefined | void>;
+        })
+      | (Partial<Record<HTTPMethod, Handler<BunRequest<Path>, Server<WebSocketData>, HTTPResponse>>> & {
+          websocket?: never;
+          upgrade?: Handler<BunRequest<Path>, Server<WebSocketData>, Response | undefined | void>;
+        });
+
     type Routes<WebSocketData, R extends string> = {
       [Path in R]:
         | BaseRouteValue
         | Handler<BunRequest<Path>, Server<WebSocketData>, Response>
-        | Partial<Record<HTTPMethod, Handler<BunRequest<Path>, Server<WebSocketData>, Response>>>;
+        | Partial<Record<HTTPMethod, Handler<BunRequest<Path>, Server<WebSocketData>, Response>>>
+        | RouteWithWebSocket<WebSocketData, Path, Response>;
     };
 
     type RoutesWithUpgrade<WebSocketData, R extends string> = {
       [Path in R]:
         | BaseRouteValue
         | Handler<BunRequest<Path>, Server<WebSocketData>, Response | undefined | void>
-        | Partial<Record<HTTPMethod, Handler<BunRequest<Path>, Server<WebSocketData>, Response | undefined | void>>>;
+        | Partial<Record<HTTPMethod, Handler<BunRequest<Path>, Server<WebSocketData>, Response | undefined | void>>>
+        | RouteWithWebSocket<WebSocketData, Path, Response | undefined | void>;
     };
 
     type FetchOrRoutes<WebSocketData, R extends string> =

src/bun.js/api/server.zig

@@ -561,10 +561,19 @@ pub fn NewServer(protocol_enum: enum { http, https }, development_kind: enum { d
         /// So we have to store it.
         user_routes: std.ArrayListUnmanaged(UserRoute) = .{},
 
+        /// Per-route WebSocket contexts. Index is (id - 2) where id comes from app.ws()
+        /// Use Shared pointers to ensure stable addresses (ServerWebSocket stores raw pointers to handlers)
+        /// When ref count reaches 0, WebSocketServerContext.deinit() is automatically called to unprotect JSValues
+        route_websocket_contexts: std.ArrayListUnmanaged(SharedWebSocketContext) = .{},
+
         on_clienterror: jsc.Strong.Optional = .empty,
 
         inspector_server_id: jsc.Debugger.DebuggerId = .init(0),
 
+        /// Shared pointer type for route-specific WebSocket contexts
+        /// .deinit = true enables automatic cleanup: when the last reference is released,
+        /// WebSocketServerContext.deinit() is called to unprotect JSValues
+        pub const SharedWebSocketContext = bun.ptr.shared.WithOptions(*WebSocketServerContext, .{ .deinit = true });
         pub const doStop = host_fn.wrapInstanceMethod(ThisServer, "stopFromJS", false);
         pub const dispose = host_fn.wrapInstanceMethod(ThisServer, "disposeFromJS", false);
         pub const doUpgrade = host_fn.wrapInstanceMethod(ThisServer, "onUpgrade", false);
@@ -578,6 +587,8 @@ pub fn NewServer(protocol_enum: enum { http, https }, development_kind: enum { d
             id: u32,
             server: *ThisServer,
             route: ServerConfig.RouteDeclaration,
+            /// Index into route_websocket_contexts, or null if no route-specific websocket
+            websocket_context_index: ?u32 = null,
 
             pub fn deinit(this: *UserRoute) void {
                 this.route.deinit();
@@ -737,8 +748,10 @@ pub fn NewServer(protocol_enum: enum { http, https }, development_kind: enum { d
         }
 
         pub fn onUpgrade(this: *ThisServer, globalThis: *jsc.JSGlobalObject, object: jsc.JSValue, optional: ?JSValue) bun.JSError!JSValue {
-            if (this.config.websocket == null) {
-                return globalThis.throwInvalidArguments("To enable websocket support, set the \"websocket\" object in Bun.serve({})", .{});
+            // Check if we have either a global websocket or route-specific websockets
+            const has_websocket = this.config.websocket != null or this.route_websocket_contexts.items.len > 0;
+            if (!has_websocket) {
+                return globalThis.throwInvalidArguments("To enable websocket support, set the \"websocket\" object in Bun.serve({}) or in a route", .{});
             }
 
             if (this.flags.terminated) {
@@ -963,6 +976,16 @@ pub fn NewServer(protocol_enum: enum { http, https }, development_kind: enum { d
                 }
             }
 
+            // Verify we have a WebSocket handler for this upgrade
+            // Either route-specific or global
+            if (upgrader.route_websocket_context_index) |ws_idx| {
+                if (ws_idx >= this.route_websocket_contexts.items.len) {
+                    return globalThis.throwInvalidArguments("Invalid WebSocket context index for this route", .{});
+                }
+            } else if (this.config.websocket == null) {
+                return globalThis.throwInvalidArguments("No WebSocket handler available for this route", .{});
+            }
+
             // Write status, custom headers, and cookies in one place
             if (fetch_headers_to_use != null or cookies_to_write != null) {
                 // we must write the status first so that 200 OK isn't written
@@ -989,7 +1012,12 @@ pub fn NewServer(protocol_enum: enum { http, https }, development_kind: enum { d
             upgrader.request_weakref.deref();
 
             data_value.ensureStillAlive();
-            const ws = ServerWebSocket.init(&this.config.websocket.?.handler, data_value, signal);
+
+            // Create ServerWebSocket with route-specific or global handler
+            const ws = if (upgrader.route_websocket_context_index) |ws_idx|
+                ServerWebSocket.initWithSharedContext(this.route_websocket_contexts.items[ws_idx], data_value, signal)
+            else
+                ServerWebSocket.init(&this.config.websocket.?.handler, data_value, signal);
             data_value.ensureStillAlive();
 
             var sec_websocket_protocol_str = sec_websocket_protocol.toSlice(bun.default_allocator);
@@ -1609,6 +1637,12 @@ pub fn NewServer(protocol_enum: enum { http, https }, development_kind: enum { d
             }
             this.user_routes.deinit(bun.default_allocator);
 
+            // Clean up route-specific WebSocket contexts
+            for (this.route_websocket_contexts.items) |*shared_ws| {
+                shared_ws.deinit(); // Decrements ref count, calls WebSocketServerContext.deinit() when count reaches 0
+            }
+            this.route_websocket_contexts.deinit(bun.default_allocator);
+
             this.config.deinit();
 
             this.on_clienterror.deinit();
@@ -2306,6 +2340,10 @@ pub fn NewServer(protocol_enum: enum { http, https }, development_kind: enum { d
             var should_deinit_context = false;
             var prepared = server.prepareJsRequestContext(req, resp, &should_deinit_context, .no, method) orelse return;
             prepared.ctx.upgrade_context = upgrade_ctx; // set the upgrade context
+
+            // Store route-specific WebSocket context index if present
+            prepared.ctx.route_websocket_context_index = this.websocket_context_index;
+
             const server_request_list = js.routeListGetCached(server.jsValueAssertAlive()).?;
             const response_value = bun.jsc.fromJSHostCall(server.globalThis, @src(), Bun__ServerRouteList__callRoute, .{ server.globalThis, index, prepared.request_object, server.jsValueAssertAlive(), server_request_list, &prepared.js_request, req }) catch |err| server.globalThis.takeException(err);
 
@@ -2314,8 +2352,10 @@ pub fn NewServer(protocol_enum: enum { http, https }, development_kind: enum { d
 
         pub fn onWebSocketUpgrade(this: *ThisServer, resp: *App.Response, req: *uws.Request, upgrade_ctx: *uws.SocketContext, id: usize) void {
             jsc.markBinding(@src());
-            if (id == 1) {
-                // This is actually a UserRoute if id is 1 so it's safe to cast
+            if (id >= 1) {
+                // This is actually a UserRoute if id >= 1 so it's safe to cast
+                // id == 1: global websocket
+                // id >= 2: route-specific websocket (context index stored in UserRoute)
                 upgradeWebSocketUserRoute(@ptrCast(this), resp, req, upgrade_ctx, null);
                 return;
             }
@@ -2481,7 +2521,23 @@ pub fn NewServer(protocol_enum: enum { http, https }, development_kind: enum { d
                     for (old_user_routes.items) |*r| r.route.deinit();
                     old_user_routes.deinit(bun.default_allocator);
                 }
+
+                // Clean up old route-specific WebSocket contexts
+                var old_route_websocket_contexts = this.route_websocket_contexts;
+                defer {
+                    // Deinit Shared pointers - this loop decrements ref counts.
+                    // With .deinit = true, when ref count reaches 0, WebSocketServerContext.deinit()
+                    // is called automatically to unprotect JSValues.
+                    for (old_route_websocket_contexts.items) |*shared_ws| {
+                        shared_ws.deinit();
+                    }
+                    // Free the slice container using bun.default_allocator
+                    old_route_websocket_contexts.deinit(bun.default_allocator);
+                }
+
                 this.user_routes = std.ArrayListUnmanaged(UserRoute).initCapacity(bun.default_allocator, user_routes_to_build_list.items.len) catch @panic("OOM");
+                this.route_websocket_contexts = std.ArrayListUnmanaged(SharedWebSocketContext).initCapacity(bun.default_allocator, user_routes_to_build_list.items.len) catch @panic("OOM");
+
                 const paths_zig = bun.default_allocator.alloc(ZigString, user_routes_to_build_list.items.len) catch @panic("OOM");
                 defer bun.default_allocator.free(paths_zig);
                 const callbacks_js = bun.default_allocator.alloc(jsc.JSValue, user_routes_to_build_list.items.len) catch @panic("OOM");
@@ -2490,10 +2546,25 @@ pub fn NewServer(protocol_enum: enum { http, https }, development_kind: enum { d
                 for (user_routes_to_build_list.items, paths_zig, callbacks_js, 0..) |*builder, *p_zig, *cb_js, i| {
                     p_zig.* = ZigString.init(builder.route.path);
                     cb_js.* = builder.callback.get().?;
+
+                    // Store route-specific WebSocket context if present
+                    var ws_ctx_index: ?u32 = null;
+                    if (builder.websocket) |ws| {
+                        ws_ctx_index = @truncate(this.route_websocket_contexts.items.len);
+                        // Use Shared pointer to ensure stable memory address for raw pointers in ServerWebSocket
+                        // .deinit = true enables automatic cleanup: when ref count reaches 0,
+                        // WebSocketServerContext.deinit() is called automatically to unprotect JSValues
+                        const shared_ws = SharedWebSocketContext.new(ws);
+                        shared_ws.get().protect();
+                        this.route_websocket_contexts.appendAssumeCapacity(shared_ws);
+                        builder.websocket = null; // Mark as moved
+                    }
+
                     this.user_routes.appendAssumeCapacity(.{
                         .id = @truncate(i),
                         .server = this,
                         .route = builder.route,
+                        .websocket_context_index = ws_ctx_index,
                     });
                     builder.route = .{}; // Mark as moved
                 }
@@ -2509,6 +2580,14 @@ pub fn NewServer(protocol_enum: enum { http, https }, development_kind: enum { d
                 websocket.handler.flags.ssl = ssl_enabled;
             }
 
+            // Setup route-specific WebSocket contexts
+            for (this.route_websocket_contexts.items) |*shared_ws| {
+                const websocket = shared_ws.get();
+                websocket.globalObject = this.globalThis;
+                websocket.handler.app = app;
+                websocket.handler.flags.ssl = ssl_enabled;
+            }
+
             // --- 3. Register compiled user routes (this.user_routes) & Track "/*" Coverage ---
             var star_methods_covered_by_user = bun.http.Method.Set.initEmpty();
             var has_any_user_route_for_star_path = false; // True if "/*" path appears in user_routes at all
@@ -2534,14 +2613,28 @@ pub fn NewServer(protocol_enum: enum { http, https }, development_kind: enum { d
                             star_methods_covered_by_user = .initFull();
                         }
 
-                        if (this.config.websocket) |*websocket| {
+                        // Register WebSocket route - prefer route-specific context over global
+                        if (user_route.websocket_context_index) |ws_idx| {
+                            // Route has its own WebSocket handler
+                            if (is_star_path) {
+                                has_any_ws_route_for_star_path = true;
+                            }
+                            const ws_context = this.route_websocket_contexts.items[ws_idx].get();
+                            app.ws(
+                                user_route.route.path,
+                                user_route,
+                                2 + ws_idx, // id = 2 + index for route-specific handlers
+                                ServerWebSocket.behavior(ThisServer, ssl_enabled, ws_context.toBehavior()),
+                            );
+                        } else if (this.config.websocket) |*websocket| {
+                            // Use global WebSocket handler
                             if (is_star_path) {
                                 has_any_ws_route_for_star_path = true;
                             }
                             app.ws(
                                 user_route.route.path,
                                 user_route,
-                                1, // id 1 means is a user route
+                                1, // id 1 means is a user route with global websocket
                                 ServerWebSocket.behavior(ThisServer, ssl_enabled, websocket.toBehavior()),
                             );
                         }
@@ -2553,13 +2646,23 @@ pub fn NewServer(protocol_enum: enum { http, https }, development_kind: enum { d
                         }
 
                         // Setup user websocket in the route if needed.
-                        if (this.config.websocket) |*websocket| {
-                            // Websocket upgrade is a GET request
-                            if (method_val == .GET) {
+                        // WebSocket upgrade is a GET request, so only register for GET or ANY methods
+                        if (method_val == .GET) {
+                            if (user_route.websocket_context_index) |ws_idx| {
+                                // Route has its own WebSocket handler
+                                const ws_context = this.route_websocket_contexts.items[ws_idx].get();
                                 app.ws(
                                     user_route.route.path,
                                     user_route,
-                                    1, // id 1 means is a user route
+                                    2 + ws_idx, // id = 2 + index for route-specific handlers
+                                    ServerWebSocket.behavior(ThisServer, ssl_enabled, ws_context.toBehavior()),
+                                );
+                            } else if (this.config.websocket) |*websocket| {
+                                // Use global WebSocket handler
+                                app.ws(
+                                    user_route.route.path,
+                                    user_route,
+                                    1, // id 1 means is a user route with global websocket
                                     ServerWebSocket.behavior(ThisServer, ssl_enabled, websocket.toBehavior()),
                                 );
                             }

src/bun.js/api/server/RequestContext.zig

@@ -44,6 +44,8 @@ pub fn NewRequestContext(comptime ssl_enabled: bool, comptime debug_mode: bool,
         flags: NewFlags(debug_mode) = .{},
 
         upgrade_context: ?*uws.SocketContext = null,
+        /// Index into server.route_websocket_contexts for route-specific WebSocket handlers
+        route_websocket_context_index: ?u32 = null,
 
         /// We can only safely free once the request body promise is finalized
         /// and the response is rejected

src/bun.js/api/server/ServerConfig.zig

@@ -581,6 +581,47 @@ pub fn fromJS(
                         HTTP.Method.TRACE,
                     };
                     var found = false;
+                    var websocket_ctx: ?WebSocketServerContext = null;
+                    var upgrade_callback: jsc.Strong.Optional = .empty;
+
+                    // Check for websocket and upgrade fields
+                    if (try value.getOwn(global, "websocket")) |ws_value| {
+                        if (!ws_value.isUndefined()) {
+                            websocket_ctx = try WebSocketServerContext.onCreate(global, ws_value);
+                        }
+                    }
+
+                    if (try value.getOwn(global, "upgrade")) |upgrade_value| {
+                        if (upgrade_value.isCallable()) {
+                            upgrade_callback = .create(upgrade_value.withAsyncContextIfNeeded(global), global);
+                        }
+                    }
+
+                    // Validate: if route has websocket, it must have upgrade
+                    if (websocket_ctx != null and upgrade_callback.impl == null) {
+                        return global.throwInvalidArguments("Route has 'websocket' but missing 'upgrade' handler. Both must be specified together.", .{});
+                    }
+
+                    // If we have an upgrade handler, add it FIRST (before other method handlers)
+                    // This ensures app.ws() is registered before app.method(.GET), so WebSocket
+                    // upgrade requests are handled by app.ws(), and regular GET requests yield
+                    // and fall through to the GET handler
+                    if (upgrade_callback.impl != null) {
+                        if (!found) {
+                            try validateRouteName(global, path);
+                        }
+                        args.user_routes_to_build.append(.{
+                            .route = .{
+                                .path = bun.handleOom(bun.default_allocator.dupeZ(u8, path)),
+                                .method = .{ .specific = .GET },
+                            },
+                            .callback = upgrade_callback,
+                            .websocket = websocket_ctx, // May be null (uses global)
+                        }) catch |err| bun.handleOom(err);
+                        found = true;
+                    }
+
+                    // Process HTTP method handlers (registered after upgrade handler)
                     inline for (methods) |method| {
                         if (try value.getOwn(global, @tagName(method))) |function| {
                             if (!found) {
@@ -589,12 +630,15 @@ pub fn fromJS(
                             found = true;
 
                             if (function.isCallable()) {
+                                // Never attach websocket to method handlers
+                                // WebSocket is handled separately via the upgrade callback
                                 args.user_routes_to_build.append(.{
                                     .route = .{
                                         .path = bun.handleOom(bun.default_allocator.dupeZ(u8, path)),
                                         .method = .{ .specific = method },
                                     },
                                     .callback = .create(function.withAsyncContextIfNeeded(global), global),
+                                    .websocket = null,
                                 }) catch |err| bun.handleOom(err);
                             } else if (try AnyRoute.fromJS(global, path, function, init_ctx)) |html_route| {
                                 var method_set = bun.http.Method.Set.initEmpty();
@@ -1070,10 +1114,14 @@ pub fn fromJS(
 const UserRouteBuilder = struct {
     route: ServerConfig.RouteDeclaration,
     callback: jsc.Strong.Optional = .empty,
+    websocket: ?WebSocketServerContext = null,
 
     pub fn deinit(this: *UserRouteBuilder) void {
         this.route.deinit();
         this.callback.deinit();
+        if (this.websocket) |ws| {
+            ws.unprotect();
+        }
     }
 };
 

src/bun.js/api/server/ServerWebSocket.zig

@@ -1,6 +1,9 @@
 const ServerWebSocket = @This();
 
 #handler: *WebSocketServer.Handler,
+/// Optional Shared pointer for route-specific WebSocket contexts.
+/// When set, this holds a reference to keep the context alive.
+#shared_context: ?bun.api.server.NewServer(.http, .debug).SharedWebSocketContext = null,
 #this_value: jsc.JSRef = .empty(),
 #flags: Flags = .{},
 #signal: ?*bun.webcore.AbortSignal = null,
@@ -51,6 +54,23 @@ pub fn init(handler: *WebSocketServer.Handler, data_value: jsc.JSValue, signal:
     return this;
 }
 
+/// Initialize a ServerWebSocket with a route-specific shared context.
+/// This clones the shared context to hold a reference and keep it alive.
+pub fn initWithSharedContext(shared_ctx: bun.api.server.NewServer(.http, .debug).SharedWebSocketContext, data_value: jsc.JSValue, signal: ?*bun.webcore.AbortSignal) *ServerWebSocket {
+    const handler = shared_ctx.get();
+    const globalObject = handler.globalObject;
+    const this = ServerWebSocket.new(.{
+        .#handler = &handler.handler,
+        .#shared_context = shared_ctx.clone(), // Clone to increment ref count
+        .#signal = signal,
+    });
+    // Get a strong ref and downgrade when terminating/close and GC will be able to collect the newly created value
+    const this_value = this.toJS(globalObject);
+    this.#this_value = .initStrong(this_value, globalObject);
+    js.dataSetCached(this_value, globalObject, data_value);
+    return this;
+}
+
 pub fn memoryCost(this: *const ServerWebSocket) usize {
     if (this.#flags.closed) {
         return @sizeOf(ServerWebSocket);
@@ -288,11 +308,7 @@ pub fn onClose(this: *ServerWebSocket, _: uws.AnyWebSocket, code: i32, message:
     var handler = this.#handler;
     const was_closed = this.isClosed();
     this.#flags.closed = true;
-    defer {
-        if (!was_closed) {
-            handler.active_connections -|= 1;
-        }
-    }
+
     const signal = this.#signal;
     this.#signal = null;
 
@@ -305,6 +321,21 @@ pub fn onClose(this: *ServerWebSocket, _: uws.AnyWebSocket, code: i32, message:
         if (this.#this_value.isNotEmpty()) {
             this.#this_value.downgrade();
         }
+
+        // Decrement active connections BEFORE releasing shared context
+        // to avoid use-after-free (handler pointer may be inside the context)
+        if (!was_closed) {
+            handler.active_connections -|= 1;
+        }
+
+        // Release the shared context reference if we have one
+        // When the last reference is released, WebSocketServerContext.deinit()
+        // will be called automatically to unprotect JSValues
+        if (this.#shared_context) |shared_ctx| {
+            var ctx = shared_ctx;
+            this.#shared_context = null;
+            ctx.deinit();
+        }
     }
 
     const vm = handler.vm;

src/bun.js/api/server/WebSocketServerContext.zig

@@ -124,6 +124,9 @@ pub fn protect(this: WebSocketServerContext) void {
 pub fn unprotect(this: WebSocketServerContext) void {
     this.handler.unprotect();
 }
+pub fn deinit(this: *WebSocketServerContext) void {
+    this.unprotect();
+}
 
 const CompressTable = bun.ComptimeStringMap(i32, .{
     .{ "disable", 0 },

test/js/bun/http/serve-route-websocket.test.ts

@@ -0,0 +1,868 @@
+import { describe, expect, test } from "bun:test";
+
+describe("Bun.serve() route-specific WebSocket handlers", () => {
+  test("route-specific websocket handlers work independently", async () => {
+    using server = Bun.serve({
+      port: 0,
+      routes: {
+        "/api/v1/chat": {
+          websocket: {
+            open(ws) {
+              ws.send("chat:welcome");
+            },
+            message(ws, data) {
+              ws.send("chat:" + data);
+            },
+          },
+          upgrade(req, server) {
+            return server.upgrade(req);
+          },
+        },
+        "/api/v2/notifications": {
+          websocket: {
+            open(ws) {
+              ws.send("notif:connected");
+            },
+            message(ws, data) {
+              ws.send("notif:" + data);
+            },
+          },
+          upgrade(req, server) {
+            return server.upgrade(req);
+          },
+        },
+      },
+    });
+
+    // Test chat WebSocket
+    const chatWs = new WebSocket(`ws://localhost:${server.port}/api/v1/chat`);
+    const chatMessages: string[] = [];
+
+    const { promise: chatResponse, resolve: resolveChatResponse } = Promise.withResolvers<void>();
+    let chatResponseCount = 0;
+    chatWs.onmessage = e => {
+      chatMessages.push(e.data);
+      chatResponseCount++;
+      if (chatResponseCount === 2) resolveChatResponse();
+    };
+
+    await new Promise(resolve => (chatWs.onopen = resolve));
+    expect(chatMessages[0]).toBe("chat:welcome");
+
+    chatWs.send("hello");
+    await chatResponse;
+    expect(chatMessages[1]).toBe("chat:hello");
+
+    chatWs.close();
+
+    // Test notifications WebSocket
+    const notifWs = new WebSocket(`ws://localhost:${server.port}/api/v2/notifications`);
+    const notifMessages: string[] = [];
+
+    const { promise: notifResponse, resolve: resolveNotifResponse } = Promise.withResolvers<void>();
+    let notifResponseCount = 0;
+    notifWs.onmessage = e => {
+      notifMessages.push(e.data);
+      notifResponseCount++;
+      if (notifResponseCount === 2) resolveNotifResponse();
+    };
+
+    await new Promise(resolve => (notifWs.onopen = resolve));
+    expect(notifMessages[0]).toBe("notif:connected");
+
+    notifWs.send("test");
+    await notifResponse;
+    expect(notifMessages[1]).toBe("notif:test");
+
+    notifWs.close();
+  });
+
+  test("route-specific websocket with data in upgrade", async () => {
+    using server = Bun.serve({
+      port: 0,
+      routes: {
+        "/ws": {
+          websocket: {
+            open(ws) {
+              ws.send("data:" + JSON.stringify(ws.data));
+            },
+          },
+          upgrade(req, server) {
+            return server.upgrade(req, {
+              data: { user: "alice", room: "general" },
+            });
+          },
+        },
+      },
+    });
+
+    const ws = new WebSocket(`ws://localhost:${server.port}/ws`);
+    const messages: string[] = [];
+    ws.onmessage = e => messages.push(e.data);
+    await new Promise(resolve => (ws.onopen = resolve));
+
+    expect(messages[0]).toBe('data:{"user":"alice","room":"general"}');
+    ws.close();
+  });
+
+  test("route-specific websocket with close handler", async () => {
+    let closeCode = 0;
+
+    const { promise: closeHandlerCalled, resolve: resolveCloseHandler } = Promise.withResolvers<void>();
+
+    using server = Bun.serve({
+      port: 0,
+      routes: {
+        "/ws": {
+          websocket: {
+            open(ws) {
+              ws.send("ready");
+            },
+            close(ws, code) {
+              closeCode = code;
+              resolveCloseHandler();
+            },
+          },
+          upgrade(req, server) {
+            return server.upgrade(req);
+          },
+        },
+      },
+    });
+
+    const ws = new WebSocket(`ws://localhost:${server.port}/ws`);
+    await new Promise(resolve => (ws.onopen = resolve));
+    ws.close(1000);
+
+    await closeHandlerCalled;
+    expect(closeCode).toBe(1000);
+  });
+
+  test("global websocket handler still works", async () => {
+    using server = Bun.serve({
+      port: 0,
+      websocket: {
+        open(ws) {
+          ws.send("global:welcome");
+        },
+        message(ws, data) {
+          ws.send("global:" + data);
+        },
+      },
+      routes: {
+        "/api/test": {
+          upgrade(req, server) {
+            return server.upgrade(req);
+          },
+        },
+      },
+    });
+
+    const ws = new WebSocket(`ws://localhost:${server.port}/api/test`);
+    const messages: string[] = [];
+
+    const { promise: messageReceived, resolve: resolveMessageReceived } = Promise.withResolvers<void>();
+    ws.onmessage = e => {
+      messages.push(e.data);
+      if (messages.length > 1) resolveMessageReceived();
+    };
+
+    await new Promise(resolve => (ws.onopen = resolve));
+
+    expect(messages[0]).toBe("global:welcome");
+
+    ws.send("test");
+    await messageReceived;
+    expect(messages[1]).toBe("global:test");
+
+    ws.close();
+  });
+
+  test("mix of route-specific and global websocket handlers", async () => {
+    using server = Bun.serve({
+      port: 0,
+      websocket: {
+        open(ws) {
+          ws.send("global:open");
+        },
+        message(ws, data) {
+          ws.send("global:" + data);
+        },
+      },
+      routes: {
+        "/specific": {
+          websocket: {
+            open(ws) {
+              ws.send("specific:open");
+            },
+            message(ws, data) {
+              ws.send("specific:" + data);
+            },
+          },
+          upgrade(req, server) {
+            return server.upgrade(req);
+          },
+        },
+        "/global": {
+          upgrade(req, server) {
+            return server.upgrade(req);
+          },
+        },
+      },
+    });
+
+    // Test route-specific handler
+    const specificWs = new WebSocket(`ws://localhost:${server.port}/specific`);
+    const specificMessages: string[] = [];
+
+    const { promise: specificMessageReceived, resolve: resolveSpecificMessage } = Promise.withResolvers<void>();
+    specificWs.onmessage = e => {
+      specificMessages.push(e.data);
+      if (specificMessages.length > 1) resolveSpecificMessage();
+    };
+
+    await new Promise(resolve => (specificWs.onopen = resolve));
+
+    expect(specificMessages[0]).toBe("specific:open");
+    specificWs.send("hello");
+    await specificMessageReceived;
+    expect(specificMessages[1]).toBe("specific:hello");
+    specificWs.close();
+
+    // Test global handler
+    const globalWs = new WebSocket(`ws://localhost:${server.port}/global`);
+    const globalMessages: string[] = [];
+
+    const { promise: globalMessageReceived, resolve: resolveGlobalMessage } = Promise.withResolvers<void>();
+    globalWs.onmessage = e => {
+      globalMessages.push(e.data);
+      if (globalMessages.length > 1) resolveGlobalMessage();
+    };
+
+    await new Promise(resolve => (globalWs.onopen = resolve));
+
+    expect(globalMessages[0]).toBe("global:open");
+    globalWs.send("world");
+    await globalMessageReceived;
+    expect(globalMessages[1]).toBe("global:world");
+    globalWs.close();
+  });
+
+  test("route-specific websocket with multiple HTTP methods", async () => {
+    let wsMessageReceived = "";
+
+    const { promise: messageProcessed, resolve: resolveMessageProcessed } = Promise.withResolvers<void>();
+
+    using server = Bun.serve({
+      port: 0,
+      routes: {
+        "/api/resource": {
+          GET() {
+            return new Response("GET response");
+          },
+          POST() {
+            return new Response("POST response");
+          },
+          websocket: {
+            open(ws) {
+              ws.send("ws:ready");
+            },
+            message(ws, data) {
+              wsMessageReceived = data.toString();
+              ws.send("ws:received");
+              resolveMessageProcessed();
+            },
+          },
+          upgrade(req, server) {
+            return server.upgrade(req);
+          },
+        },
+      },
+    });
+
+    // Test HTTP GET
+    const getResp = await fetch(`http://localhost:${server.port}/api/resource`);
+    expect(await getResp.text()).toBe("GET response");
+
+    // Test HTTP POST
+    const postResp = await fetch(`http://localhost:${server.port}/api/resource`, { method: "POST" });
+    expect(await postResp.text()).toBe("POST response");
+
+    // Test WebSocket (which uses GET under the hood)
+    const ws = new WebSocket(`ws://localhost:${server.port}/api/resource`);
+    const messages: string[] = [];
+
+    const { promise: messageReceived, resolve: resolveMessageReceived } = Promise.withResolvers<void>();
+    ws.onmessage = e => {
+      messages.push(e.data);
+      if (messages.length > 1) resolveMessageReceived();
+    };
+
+    await new Promise(resolve => (ws.onopen = resolve));
+
+    expect(messages[0]).toBe("ws:ready");
+    ws.send("test-message");
+    await Promise.all([messageReceived, messageProcessed]);
+    expect(messages[1]).toBe("ws:received");
+    expect(wsMessageReceived).toBe("test-message");
+    ws.close();
+  });
+
+  test("route-specific websocket without upgrade handler errors appropriately", () => {
+    // Should throw an error because websocket requires upgrade handler
+    expect(() => {
+      Bun.serve({
+        port: 0,
+        routes: {
+          "/ws": {
+            websocket: {
+              open(ws) {
+                ws.send("should not reach here");
+              },
+            },
+            // Note: no upgrade handler
+            GET() {
+              return new Response("This is not a WebSocket endpoint");
+            },
+          },
+        },
+      });
+    }).toThrow("Route has 'websocket' but missing 'upgrade' handler");
+  });
+
+  test("server.reload() preserves route-specific websocket handlers", async () => {
+    const { promise, resolve } = Promise.withResolvers<void>();
+    let stage = 0;
+
+    using server = Bun.serve({
+      port: 0,
+      routes: {
+        "/ws": {
+          websocket: {
+            open(ws) {
+              ws.send(`stage${stage}:open`);
+            },
+          },
+          upgrade(req, server) {
+            return server.upgrade(req);
+          },
+        },
+      },
+      async fetch(req, server) {
+        if (req.url.endsWith("/reload")) {
+          stage = 1;
+          server.reload({
+            routes: {
+              "/ws": {
+                websocket: {
+                  open(ws) {
+                    ws.send("reloaded:open");
+                  },
+                },
+                upgrade(req, server) {
+                  return server.upgrade(req);
+                },
+              },
+            },
+          });
+          resolve();
+          return new Response("reloaded");
+        }
+        return new Response("not found", { status: 404 });
+      },
+    });
+
+    // Connect before reload
+    const ws1 = new WebSocket(`ws://localhost:${server.port}/ws`);
+    const messages1: string[] = [];
+    ws1.onmessage = e => messages1.push(e.data);
+    await new Promise(resolve => (ws1.onopen = resolve));
+    expect(messages1[0]).toBe("stage0:open");
+    ws1.close();
+
+    // Trigger reload
+    await fetch(`http://localhost:${server.port}/reload`);
+    await promise;
+
+    // Connect after reload
+    const ws2 = new WebSocket(`ws://localhost:${server.port}/ws`);
+    const messages2: string[] = [];
+    ws2.onmessage = e => messages2.push(e.data);
+    await new Promise(resolve => (ws2.onopen = resolve));
+    expect(messages2[0]).toBe("reloaded:open");
+    ws2.close();
+  });
+
+  test("server.reload() removes websocket handler", async () => {
+    using server = Bun.serve({
+      port: 0,
+      routes: {
+        "/ws": {
+          websocket: {
+            open(ws) {
+              ws.send("initial");
+            },
+          },
+          upgrade(req, server) {
+            return server.upgrade(req);
+          },
+        },
+      },
+    });
+
+    // Connect with websocket handler
+    const ws1 = new WebSocket(`ws://localhost:${server.port}/ws`);
+    const messages1: string[] = [];
+    ws1.onmessage = e => messages1.push(e.data);
+    await new Promise(resolve => (ws1.onopen = resolve));
+    expect(messages1[0]).toBe("initial");
+    ws1.close();
+
+    // Reload without websocket handler
+    server.reload({
+      routes: {
+        "/ws": {
+          GET() {
+            return new Response("no websocket");
+          },
+        },
+      },
+    });
+
+    // Regular GET should work
+    const resp = await fetch(`http://localhost:${server.port}/ws`);
+    expect(await resp.text()).toBe("no websocket");
+
+    // WebSocket should fail
+    const ws2 = new WebSocket(`ws://localhost:${server.port}/ws`);
+    const { promise: errorOccurred, resolve: resolveError } = Promise.withResolvers<void>();
+    ws2.onerror = () => {
+      resolveError();
+    };
+    await errorOccurred;
+  });
+
+  test("server.reload() adds websocket handler to existing route", async () => {
+    using server = Bun.serve({
+      port: 0,
+      routes: {
+        "/ws": {
+          GET() {
+            return new Response("no websocket yet");
+          },
+        },
+      },
+    });
+
+    // Regular GET should work
+    const resp1 = await fetch(`http://localhost:${server.port}/ws`);
+    expect(await resp1.text()).toBe("no websocket yet");
+
+    // Reload with websocket handler
+    server.reload({
+      routes: {
+        "/ws": {
+          GET() {
+            return new Response("now has websocket");
+          },
+          websocket: {
+            open(ws) {
+              ws.send("added");
+            },
+          },
+          upgrade(req, server) {
+            return server.upgrade(req);
+          },
+        },
+      },
+    });
+
+    // Regular GET should still work
+    const resp2 = await fetch(`http://localhost:${server.port}/ws`);
+    expect(await resp2.text()).toBe("now has websocket");
+
+    // WebSocket should now work
+    const ws = new WebSocket(`ws://localhost:${server.port}/ws`);
+    const messages: string[] = [];
+    ws.onmessage = e => messages.push(e.data);
+    await new Promise(resolve => (ws.onopen = resolve));
+    expect(messages[0]).toBe("added");
+    ws.close();
+  });
+
+  test("server.reload() with active websocket connections", async () => {
+    let messageReceived = "";
+
+    const { promise: messageProcessed, resolve: resolveMessageProcessed } = Promise.withResolvers<void>();
+
+    using server = Bun.serve({
+      port: 0,
+      routes: {
+        "/ws": {
+          websocket: {
+            open(ws) {
+              ws.send("v1");
+            },
+            message(ws, data) {
+              messageReceived = data.toString();
+              ws.send("v1:echo");
+              resolveMessageProcessed();
+            },
+          },
+          upgrade(req, server) {
+            return server.upgrade(req);
+          },
+        },
+      },
+    });
+
+    // Create active connection
+    const ws = new WebSocket(`ws://localhost:${server.port}/ws`);
+    const messages: string[] = [];
+
+    const { promise: messageReceived1, resolve: resolveMessageReceived1 } = Promise.withResolvers<void>();
+    ws.onmessage = e => {
+      messages.push(e.data);
+      if (messages.length > 1) resolveMessageReceived1();
+    };
+
+    await new Promise(resolve => (ws.onopen = resolve));
+    expect(messages[0]).toBe("v1");
+
+    // Reload while connection is active
+    server.reload({
+      routes: {
+        "/ws": {
+          websocket: {
+            open(ws) {
+              ws.send("v2");
+            },
+            message(ws, data) {
+              ws.send("v2:echo");
+            },
+          },
+          upgrade(req, server) {
+            return server.upgrade(req);
+          },
+        },
+      },
+    });
+
+    // Existing connection should still use old handlers
+    ws.send("test");
+    await Promise.all([messageReceived1, messageProcessed]);
+    expect(messages[1]).toBe("v1:echo");
+    expect(messageReceived).toBe("test");
+    ws.close();
+
+    // New connection should use new handlers
+    const ws2 = new WebSocket(`ws://localhost:${server.port}/ws`);
+    const messages2: string[] = [];
+
+    const { promise: messageReceived2, resolve: resolveMessageReceived2 } = Promise.withResolvers<void>();
+    ws2.onmessage = e => {
+      messages2.push(e.data);
+      if (messages2.length > 1) resolveMessageReceived2();
+    };
+
+    await new Promise(resolve => (ws2.onopen = resolve));
+    expect(messages2[0]).toBe("v2");
+    ws2.send("test2");
+    await messageReceived2;
+    expect(messages2[1]).toBe("v2:echo");
+    ws2.close();
+  });
+
+  test("multiple concurrent websocket connections to same route", async () => {
+    const openCount = { count: 0 };
+    const messageCount = { count: 0 };
+
+    const { promise: allMessagesReceived, resolve: resolveAllMessagesReceived } = Promise.withResolvers<void>();
+
+    using server = Bun.serve({
+      port: 0,
+      routes: {
+        "/ws": {
+          websocket: {
+            open(ws) {
+              openCount.count++;
+              ws.send(`connection-${openCount.count}`);
+            },
+            message(ws, data) {
+              messageCount.count++;
+              ws.send(`echo-${data}`);
+              if (messageCount.count === 5) resolveAllMessagesReceived();
+            },
+          },
+          upgrade(req, server) {
+            return server.upgrade(req);
+          },
+        },
+      },
+    });
+
+    // Create 5 concurrent connections with promise resolvers for each
+    const connectionPromises = Array.from({ length: 5 }, (_, i) => {
+      const { promise, resolve } = Promise.withResolvers<void>();
+      return { promise, resolve, id: i };
+    });
+
+    const connections = await Promise.all(
+      connectionPromises.map(async ({ promise, resolve, id }) => {
+        const ws = new WebSocket(`ws://localhost:${server.port}/ws`);
+        const messages: string[] = [];
+        ws.onmessage = e => {
+          messages.push(e.data);
+          if (messages.length >= 2) resolve();
+        };
+        await new Promise(resolveOpen => (ws.onopen = resolveOpen));
+        return { ws, messages, id, promise };
+      }),
+    );
+
+    expect(openCount.count).toBe(5);
+
+    // Each should have unique connection message
+    for (let i = 0; i < 5; i++) {
+      expect(connections[i].messages[0]).toMatch(/^connection-\d+$/);
+    }
+
+    // Send messages from all connections
+    for (const conn of connections) {
+      conn.ws.send(`msg-${conn.id}`);
+    }
+
+    // Wait for server to receive all messages
+    await allMessagesReceived;
+    expect(messageCount.count).toBe(5);
+
+    // Wait for all echo responses to arrive back at clients
+    await Promise.all(connections.map(conn => conn.promise));
+
+    // Each should get their echo back
+    for (const conn of connections) {
+      expect(conn.messages[1]).toBe(`echo-msg-${conn.id}`);
+      conn.ws.close();
+    }
+  });
+
+  test("multiple concurrent websocket connections to different routes", async () => {
+    using server = Bun.serve({
+      port: 0,
+      routes: {
+        "/chat": {
+          websocket: {
+            open(ws) {
+              ws.send("chat");
+            },
+          },
+          upgrade(req, server) {
+            return server.upgrade(req);
+          },
+        },
+        "/notifications": {
+          websocket: {
+            open(ws) {
+              ws.send("notif");
+            },
+          },
+          upgrade(req, server) {
+            return server.upgrade(req);
+          },
+        },
+        "/updates": {
+          websocket: {
+            open(ws) {
+              ws.send("updates");
+            },
+          },
+          upgrade(req, server) {
+            return server.upgrade(req);
+          },
+        },
+      },
+    });
+
+    // Connect to all routes simultaneously
+    const [chat, notif, updates] = await Promise.all([
+      (async () => {
+        const ws = new WebSocket(`ws://localhost:${server.port}/chat`);
+        const messages: string[] = [];
+        ws.onmessage = e => messages.push(e.data);
+        await new Promise(resolve => (ws.onopen = resolve));
+        return { ws, messages };
+      })(),
+      (async () => {
+        const ws = new WebSocket(`ws://localhost:${server.port}/notifications`);
+        const messages: string[] = [];
+        ws.onmessage = e => messages.push(e.data);
+        await new Promise(resolve => (ws.onopen = resolve));
+        return { ws, messages };
+      })(),
+      (async () => {
+        const ws = new WebSocket(`ws://localhost:${server.port}/updates`);
+        const messages: string[] = [];
+        ws.onmessage = e => messages.push(e.data);
+        await new Promise(resolve => (ws.onopen = resolve));
+        return { ws, messages };
+      })(),
+    ]);
+
+    expect(chat.messages[0]).toBe("chat");
+    expect(notif.messages[0]).toBe("notif");
+    expect(updates.messages[0]).toBe("updates");
+
+    chat.ws.close();
+    notif.ws.close();
+    updates.ws.close();
+  });
+
+  test("websocket with only open handler (no message/close)", async () => {
+    using server = Bun.serve({
+      port: 0,
+      routes: {
+        "/ws": {
+          websocket: {
+            open(ws) {
+              ws.send("opened");
+            },
+            // No message or close handlers
+          },
+          upgrade(req, server) {
+            return server.upgrade(req);
+          },
+        },
+      },
+    });
+
+    const ws = new WebSocket(`ws://localhost:${server.port}/ws`);
+    const messages: string[] = [];
+    ws.onmessage = e => messages.push(e.data);
+    await new Promise(resolve => (ws.onopen = resolve));
+
+    expect(messages[0]).toBe("opened");
+
+    // Should be able to send messages even without handler
+    ws.send("test");
+
+    // Should be able to close
+    ws.close();
+  });
+
+  test("websocket error handler is called on server-side exceptions", async () => {
+    const { promise: errorHandlerCalled, resolve: resolveErrorHandler } = Promise.withResolvers<void>();
+
+    using server = Bun.serve({
+      port: 0,
+      routes: {
+        "/ws": {
+          websocket: {
+            message(ws, message) {
+              // Trigger an error when receiving "trigger-error"
+              if (message === "trigger-error") {
+                throw new Error("Intentional test error");
+              }
+            },
+            error(ws, error) {
+              resolveErrorHandler();
+            },
+          },
+          upgrade(req, server) {
+            return server.upgrade(req);
+          },
+        },
+      },
+    });
+
+    const ws = new WebSocket(`ws://localhost:${server.port}/ws`);
+    await new Promise(resolve => (ws.onopen = resolve));
+
+    // Send message that triggers server-side error
+    ws.send("trigger-error");
+
+    // Wait for error handler to be called
+    await errorHandlerCalled;
+
+    ws.close();
+  });
+
+  test("server.stop() with active websocket connections", async () => {
+    const server = Bun.serve({
+      port: 0,
+      routes: {
+        "/ws": {
+          websocket: {
+            open(ws) {
+              ws.send("connected");
+            },
+            close(ws) {
+              // Close handler is called when connection closes
+            },
+          },
+          upgrade(req, server) {
+            return server.upgrade(req);
+          },
+        },
+      },
+    });
+
+    const ws = new WebSocket(`ws://localhost:${server.port}/ws`);
+    await new Promise(resolve => (ws.onopen = resolve));
+
+    // Manually close the connection and wait for it
+    const closePromise = new Promise(resolve => (ws.onclose = resolve));
+    ws.close();
+    await closePromise;
+
+    // Now stop server and await completion
+    await server.stop();
+
+    // Server should stop successfully even after WebSocket was used
+    expect(server.port).toBe(0);
+  });
+
+  test("multiple routes with same path but different methods and websocket", async () => {
+    using server = Bun.serve({
+      port: 0,
+      routes: {
+        "/api": {
+          GET() {
+            return new Response("get");
+          },
+          POST() {
+            return new Response("post");
+          },
+          PUT() {
+            return new Response("put");
+          },
+          websocket: {
+            open(ws) {
+              ws.send("ws");
+            },
+          },
+          upgrade(req, server) {
+            return server.upgrade(req);
+          },
+        },
+      },
+    });
+
+    // Test all HTTP methods work
+    const getResp = await fetch(`http://localhost:${server.port}/api`);
+    expect(await getResp.text()).toBe("get");
+
+    const postResp = await fetch(`http://localhost:${server.port}/api`, { method: "POST" });
+    expect(await postResp.text()).toBe("post");
+
+    const putResp = await fetch(`http://localhost:${server.port}/api`, { method: "PUT" });
+    expect(await putResp.text()).toBe("put");
+
+    // Test WebSocket works
+    const ws = new WebSocket(`ws://localhost:${server.port}/api`);
+    const messages: string[] = [];
+    ws.onmessage = e => messages.push(e.data);
+    await new Promise(resolve => (ws.onopen = resolve));
+    expect(messages[0]).toBe("ws");
+    ws.close();
+  });
+});