feat(dev-overrides): Add automatic response cleanup via onClose callback

changeset

changeset

Author: sommeeeer

.changeset/eleven-moose-heal.md

@@ -0,0 +1,5 @@
+---
+"@opennextjs/aws": patch
+---
+
+fix(dev-overrides): Add automatic response cleanup via onClose callback

packages/open-next/src/http/openNextResponse.ts

@@ -84,6 +84,14 @@ export class OpenNextNodeResponse extends Transform implements ServerResponse {
     ) {
       this.statusCode = statusCode;
     }
+
+    // We want to destroy this response when the original response is closed. (i.e when the client disconnects)
+    // This is to support `request.signal.onabort` in route handlers
+    if (streamCreator?.onClose) {
+      streamCreator.onClose(() => {
+        this.destroy();
+      });
+    }
   }
 
   // Necessary for next 12

packages/open-next/src/overrides/wrappers/express-dev.ts

@@ -41,6 +41,8 @@ const wrapper: WrapperHandler = async (handler, converter) => {
       req.headers["x-forwarded-proto"] = req.protocol;
     }
     const internalEvent = await converter.convertFrom(req);
+
+    let onCloseCallback: (() => void) | undefined;
     const streamCreator: StreamCreator = {
       writeHeaders: (prelude) => {
         res.setHeader("Set-Cookie", prelude.cookies);
@@ -49,7 +51,17 @@ const wrapper: WrapperHandler = async (handler, converter) => {
         return res;
       },
       onFinish: () => {},
+      onClose: (callback) => {
+        onCloseCallback = callback;
+      },
     };
+
+    res.on("close", () => {
+      if (onCloseCallback) {
+        onCloseCallback();
+      }
+    });
+
     await handler(internalEvent, { streamCreator });
   });
 

packages/open-next/src/overrides/wrappers/node.ts

@@ -8,21 +8,35 @@ import { debug, error } from "../../adapters/logger";
 const wrapper: WrapperHandler = async (handler, converter) => {
   const server = createServer(async (req, res) => {
     const internalEvent = await converter.convertFrom(req);
+    
+    let onCloseCallback: (() => void) | undefined;
     const streamCreator: StreamCreator = {
       writeHeaders: (prelude) => {
         res.setHeader("Set-Cookie", prelude.cookies);
         res.writeHead(prelude.statusCode, prelude.headers);
         res.flushHeaders();
         return res;
       },
+      onClose: (callback) => {
+        onCloseCallback = callback;
+      },
     };
+
+    res.on("close", () => {
+      if (onCloseCallback) {
+        onCloseCallback();
+      }
+    });
+
     if (internalEvent.rawPath === "/__health") {
       res.writeHead(200, {
         "Content-Type": "text/plain",
       });
       res.end("OK");
     } else {
-      await handler(internalEvent, { streamCreator });
+      await handler(internalEvent, {
+        streamCreator,
+      });
     }
   });
 

packages/open-next/src/types/open-next.ts

@@ -49,6 +49,8 @@ export interface StreamCreator {
   // Just to fix an issue with aws lambda streaming with empty body
   onWrite?: () => void;
   onFinish?: (length: number) => void;
+  // This is called when the wrappers response is closed, i.e. when the client disconnects
+  onClose?: (callback: () => void) => void;
 }
 
 export type WaitUntil = (promise: Promise<void>) => void;