fix: handle 304 responses from edge functions (#364)

* test: add test for edge function returning 304 response

* fix: maintain bodyless responses

* fix: handle cancellations

Author: pieh

packages/dev/src/main.test.ts

@@ -606,6 +606,14 @@ describe('Handling requests', () => {
 
            export const config = { path: "/hello/terminate/*" };`,
         )
+        .withFile(
+          'netlify/edge-functions/bodyless-response.mjs',
+          `export default async (req, context) => {
+            return new Response(null, { status: 304 });
+          };
+
+          export const config = { path: "/bodyless-response" };`,
+        )
       const directory = await fixture.create()
 
       vi.stubEnv('SOME_ZSH_THING_MAYBE', 'value on developer machine')
@@ -671,6 +679,11 @@ describe('Handling requests', () => {
         url: req2URL.toString(),
       })
 
+      const req3 = new Request('https://site.netlify/bodyless-response')
+      const res3 = await dev.handle(req3)
+
+      expect(res3?.status).toBe(304)
+
       await dev.stop()
       await fixture.destroy()
     })

packages/edge-functions/dev/deno/invoke.mjs

@@ -14,7 +14,7 @@
  * @param {number} requestTimeout
  */
 export function invoke(req, functions, requestTimeout) {
-  return new Promise((resolve, reject) => {
+  return new Promise(async (resolve, reject) => {
     const worker = new Worker(new URL('./workers/runner.mjs', import.meta.url).href, {
       type: 'module',
     })
@@ -23,6 +23,8 @@ export function invoke(req, functions, requestTimeout) {
     let response = null
     /** @type {ReadableStreamDefaultController<Uint8Array> | null} */
     let streamController = null
+    /** @type {ReadableStream | null} */
+    let stream = null
 
     const timeoutCheck = setTimeout(() => {
       if (!response) {
@@ -34,24 +36,6 @@ export function invoke(req, functions, requestTimeout) {
       }
     }, requestTimeout)
 
-    const stream = new ReadableStream({
-      async start(controller) {
-        streamController = controller
-
-        worker.postMessage({
-          type: 'request',
-          data: {
-            body: await req.arrayBuffer(),
-            functions,
-            headers: Object.fromEntries(req.headers.entries()),
-            method: req.method,
-            timeout: requestTimeout,
-            url: req.url,
-          },
-        })
-      },
-    })
-
     worker.onmessage = (e) => {
       const message = /** @type {Message} */ (e.data)
 
@@ -63,6 +47,18 @@ export function invoke(req, functions, requestTimeout) {
         }
 
         case 'responseStart': {
+          if (message.data.hasBody) {
+            stream = new ReadableStream({
+              start(controller) {
+                streamController = controller
+              },
+              cancel() {
+                streamController = null
+                worker.terminate()
+              },
+            })
+          }
+
           response = new Response(stream, {
             headers: message.data.headers,
             status: message.data.status,
@@ -97,5 +93,17 @@ export function invoke(req, functions, requestTimeout) {
 
       reject(e)
     }
+
+    worker.postMessage({
+      type: 'request',
+      data: {
+        body: await req.arrayBuffer(),
+        functions,
+        headers: Object.fromEntries(req.headers.entries()),
+        method: req.method,
+        timeout: requestTimeout,
+        url: req.url,
+      },
+    })
   })
 }

packages/edge-functions/dev/deno/workers/runner.mjs

@@ -53,6 +53,7 @@ worker.addEventListener('message', async (e) => {
         data: {
           headers: Object.fromEntries(res.headers.entries()),
           status: res.status,
+          hasBody: Boolean(res.body),
         },
       }),
     )

packages/edge-functions/dev/deno/workers/types.ts

@@ -48,6 +48,7 @@ export interface RunResponseStartMessage {
   data: {
     headers: Record<string, string>
     status: number
+    hasBody: boolean
   }
 }