Allow returning a Response object when upgrading

2 files changed, 137 insertions, 6 deletions

diff --git a/src/bun.js/api/server.zig b/src/bun.js/api/server.zig
index 0adaf8b3b..ce0b388d7 100644
--- a/src/bun.js/api/server.zig
+++ b/src/bun.js/api/server.zig
@@ -658,6 +658,13 @@ fn NewRequestContext(comptime ssl_enabled: bool, comptime debug_mode: bool, comp
                 ctx.finalizeForAbort();
                 return;
             }
+
+            // I don't think this case happens?
+            if (ctx.didUpgradeWebSocket()) {
+                ctx.finalize();
+                return;
+            }
+
             if (!ctx.resp.hasResponded()) {
                 ctx.renderMissing();
                 return;
@@ -1438,15 +1445,16 @@ fn NewRequestContext(comptime ssl_enabled: bool, comptime debug_mode: bool, comp
                 return;
             }
 
+            // if you return a Response object or a Promise<Response>
+            // but you upgraded the connection to a WebSocket
+            // just ignore the Response object. It doesn't do anything.
+            // it's better to do that than to throw an error
             if (ctx.didUpgradeWebSocket()) {
                 ctx.finalize();
                 return;
             }
 
             if (response_value.isEmptyOrUndefinedOrNull()) {
-                if (ctx.didUpgradeWebSocket())
-                    return;
-
                 ctx.renderMissing();
                 return;
             }
@@ -1489,10 +1497,17 @@ fn NewRequestContext(comptime ssl_enabled: bool, comptime debug_mode: bool, comp
                     .Pending => {},
                     .Fulfilled => {
                         const fulfilled_value = promise.result(vm.global.vm());
-                        if (fulfilled_value.isEmptyOrUndefinedOrNull()) {
-                            if (ctx.didUpgradeWebSocket())
-                                return;
 
+                        // if you return a Response object or a Promise<Response>
+                        // but you upgraded the connection to a WebSocket
+                        // just ignore the Response object. It doesn't do anything.
+                        // it's better to do that than to throw an error
+                        if (ctx.didUpgradeWebSocket()) {
+                            ctx.finalize();
+                            return;
+                        }
+
+                        if (fulfilled_value.isEmptyOrUndefinedOrNull()) {
                             ctx.renderMissing();
                             return;
                         }
diff --git a/test/bun.js/websocket-server.test.ts b/test/bun.js/websocket-server.test.ts
index 3b9aae5b6..eed5ffdc4 100644
--- a/test/bun.js/websocket-server.test.ts
+++ b/test/bun.js/websocket-server.test.ts
@@ -155,6 +155,122 @@ describe("websocket server", () => {
     server.stop();
   });
 
+  it("fetch() allows a Response object to be returned for an upgraded ServerWebSocket", () => {
+    var server = serve({
+      port: getPort(),
+      websocket: {
+        open(ws) {},
+        message(ws, msg) {
+          ws.send("hello world");
+        },
+      },
+      error(err) {
+        console.error(err);
+      },
+      fetch(req, server) {
+        if (
+          server.upgrade(req, {
+            data: "hello world",
+
+            // check that headers works
+            headers: {
+              "x-a": "text/plain",
+            },
+          })
+        ) {
+          if (server.upgrade(req)) {
+            throw new Error("should not upgrade twice");
+          }
+          return new Response("lol!", {
+            status: 101,
+          });
+        }
+
+        return new Response("noooooo hello world");
+      },
+    });
+
+    return new Promise((resolve, reject) => {
+      const websocket = new WebSocket(`ws://${server.hostname}:${server.port}`);
+      websocket.onopen = () => {
+        websocket.send("hello world");
+      };
+      websocket.onmessage = (e) => {
+        try {
+          expect(e.data).toBe("hello world");
+          resolve();
+        } catch (r) {
+          reject(r);
+          return;
+        } finally {
+          server?.stop();
+          websocket.close();
+        }
+      };
+      websocket.onerror = (e) => {
+        reject(e);
+      };
+    });
+  });
+
+  it("fetch() allows a Promise<Response> object to be returned for an upgraded ServerWebSocket", () => {
+    var server = serve({
+      port: getPort(),
+      websocket: {
+        async open(ws) {},
+        async message(ws, msg) {
+          await 1;
+          ws.send("hello world");
+        },
+      },
+      error(err) {
+        console.error(err);
+      },
+      async fetch(req, server) {
+        await 1;
+        if (
+          server.upgrade(req, {
+            data: "hello world",
+
+            // check that headers works
+            headers: {
+              "x-a": "text/plain",
+            },
+          })
+        ) {
+          if (server.upgrade(req)) {
+            throw new Error("should not upgrade twice");
+          }
+          return new Response("lol!", {
+            status: 101,
+          });
+        }
+
+        return new Response("noooooo hello world");
+      },
+    });
+    return new Promise((resolve, reject) => {
+      const websocket = new WebSocket(`ws://${server.hostname}:${server.port}`);
+      websocket.onopen = () => {
+        websocket.send("hello world");
+      };
+      websocket.onmessage = (e) => {
+        try {
+          expect(e.data).toBe("hello world");
+          resolve();
+        } catch (r) {
+          reject(r);
+          return;
+        } finally {
+          server?.stop();
+          websocket.close();
+        }
+      };
+      websocket.onerror = (e) => {
+        reject(e);
+      };
+    });
+  });
   it("binaryType works", async () => {
     var done = false;
     var server = serve({