Fix for lambda streaming on empty body (#443)

* fix for lambda streaming on empty body

* Create little-pots-raise.md

.changeset/little-pots-raise.md

@@ -0,0 +1,5 @@
+---
+"open-next": patch
+---
+
+Fix for lambda streaming on empty body

examples/sst/stacks/AppRouter.ts

@@ -4,6 +4,9 @@ export function AppRouter({ stack }) {
   // We should probably switch to ion once it's ready
   const site = new OpenNextCdkReferenceImplementation(stack, "approuter", {
     path: "../app-router",
+    environment: {
+      OPEN_NEXT_FORCE_NON_EMPTY_RESPONSE: "true",
+    },
   });
   // const site = new NextjsSite(stack, "approuter", {
   //   path: "../app-router",

examples/sst/stacks/OpenNextReferenceImplementation.ts

@@ -110,6 +110,7 @@ interface OpenNextOutput {
 
 interface OpenNextCdkReferenceImplementationProps {
   path: string;
+  environment?: Record<string, string>;
 }
 
 export class OpenNextCdkReferenceImplementation extends Construct {
@@ -122,6 +123,8 @@ export class OpenNextCdkReferenceImplementation extends Construct {
   private staticCachePolicy: ICachePolicy;
   private serverCachePolicy: CachePolicy;
 
+  private customEnvironment: Record<string, string>;
+
   public distribution: Distribution;
 
   constructor(
@@ -130,6 +133,7 @@ export class OpenNextCdkReferenceImplementation extends Construct {
     props: OpenNextCdkReferenceImplementationProps,
   ) {
     super(scope, id);
+    this.customEnvironment = props.environment ?? {};
     this.openNextBasePath = path.join(process.cwd(), props.path);
     execSync("npm run openbuild", {
       cwd: path.join(process.cwd(), props.path),
@@ -312,6 +316,7 @@ export class OpenNextCdkReferenceImplementation extends Construct {
       // Those 2 are used only for image optimizer
       BUCKET_NAME: this.bucket.bucketName,
       BUCKET_KEY_PREFIX: "_assets",
+      ...this.customEnvironment,
     };
   }
 

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

@@ -7,6 +7,7 @@ import type {
 import { Socket } from "net";
 import { Transform, TransformCallback, Writable } from "stream";
 
+import { debug } from "../adapters/logger";
 import { parseCookies, parseHeaders } from "./util";
 
 const SET_COOKIE_HEADER = "set-cookie";
@@ -20,7 +21,7 @@ export interface StreamCreator {
   }): Writable;
   // Just to fix an issue with aws lambda streaming with empty body
   onWrite?: () => void;
-  onFinish: () => void;
+  onFinish: (length: number) => void;
 }
 
 // We only need to implement the methods that are used by next.js
@@ -91,7 +92,8 @@ export class OpenNextNodeResponse extends Transform implements ServerResponse {
         this.flushHeaders();
       }
       onEnd(this.headers);
-      this.streamCreator?.onFinish();
+      const bodyLength = this.body.length;
+      this.streamCreator?.onFinish(bodyLength);
     });
   }
 
@@ -281,4 +283,25 @@ export class OpenNextNodeResponse extends Transform implements ServerResponse {
     this._internalWrite(chunk, encoding);
     callback();
   }
+
+  //This is only here because of aws broken streaming implementation.
+  //Hopefully one day they will be able to give us a working streaming implementation in lambda for everyone
+  //If you're lucky you have a working streaming implementation in your aws account and don't need this
+  //If not you can set the OPEN_NEXT_FORCE_NON_EMPTY_RESPONSE env variable to true
+  //BE CAREFUL: Aws keeps rolling out broken streaming implementations even on accounts that had working ones before
+  //This is not dependent on the node runtime used
+  //There is another known issue with aws lambda streaming where the request reach the lambda only way after the request has been sent by the client. For this there is absolutely nothing we can do, contact aws support if that's your case
+  _flush(callback: TransformCallback): void {
+    if (
+      this.body.length < 1 &&
+      // We use an env variable here because not all aws account have the same behavior
+      // On some aws accounts the response will hang if the body is empty
+      // We are modifying the response body here, this is not a good practice
+      process.env.OPEN_NEXT_FORCE_NON_EMPTY_RESPONSE === "true"
+    ) {
+      debug('Force writing "SOMETHING" to the response body');
+      this.push("SOMETHING");
+    }
+    callback();
+  }
 }

packages/open-next/src/wrappers/aws-lambda-streaming.ts

@@ -31,7 +31,6 @@ const handler: WrapperHandler = async (handler, converter) =>
       }
 
       const internalEvent = await converter.convertFrom(event);
-      let _hasWriten = false;
 
       //Handle compression
       const acceptEncoding =
@@ -89,13 +88,9 @@ const handler: WrapperHandler = async (handler, converter) =>
           return compressedStream ?? responseStream;
         },
         onWrite: () => {
-          _hasWriten = true;
-        },
-        onFinish: () => {
-          if (!_hasWriten) {
-            compressedStream?.end(new Uint8Array(8));
-          }
+          // _hasWriten = true;
         },
+        onFinish: () => {},
       };
 
       const response = await handler(internalEvent, streamCreator);