Refactor AWS Lambda streaming support and enhance converters

packages/aws/src/overrides/converters/aws-apigw-v1.ts

@@ -3,11 +3,12 @@ import {
 	extractHostFromHeaders,
 	removeUndefinedFromQuery,
 } from "@opennextjs/core/overrides/converters/utils.js";
-import type { InternalEvent, InternalResult } from "@opennextjs/core/types/open-next.js";
+import type { InternalEvent } from "@opennextjs/core/types/open-next.js";
 import type { Converter } from "@opennextjs/core/types/overrides.js";
-import { fromReadableStream } from "@opennextjs/core/utils/stream.js";
 import type { APIGatewayProxyEvent, APIGatewayProxyResult } from "aws-lambda";
 
+import { createBufferedStreamCreator } from "./response-stream.js";
+
 function normalizeAPIGatewayProxyEventHeaders(event: APIGatewayProxyEvent): Record<string, string> {
 	event.multiValueHeaders;
 	const headers: Record<string, string> = {};
@@ -89,10 +90,14 @@ async function convertFromAPIGatewayProxyEvent(event: APIGatewayProxyEvent): Pro
 	};
 }
 
-async function convertToApiGatewayProxyResult(result: InternalResult): Promise<APIGatewayProxyResult> {
+function convertToApiGatewayProxyResult(
+	prelude: { statusCode: number; cookies: string[]; headers: Record<string, string> },
+	body: Buffer,
+	isBase64Encoded: boolean
+): APIGatewayProxyResult {
 	const headers: Record<string, string> = {};
 	const multiValueHeaders: Record<string, string[]> = {};
-	Object.entries(result.headers).forEach(([key, value]) => {
+	Object.entries(prelude.headers).forEach(([key, value]) => {
 		if (Array.isArray(value)) {
 			multiValueHeaders[key] = value;
 		} else {
@@ -103,22 +108,26 @@ async function convertToApiGatewayProxyResult(result: InternalResult): Promise<A
 			headers[key] = value;
 		}
 	});
-
-	const body = await fromReadableStream(result.body, result.isBase64Encoded);
+	if (prelude.cookies.length > 0) {
+		multiValueHeaders["set-cookie"] = prelude.cookies;
+	}
 
 	const response: APIGatewayProxyResult = {
-		statusCode: result.statusCode,
+		statusCode: prelude.statusCode,
 		headers,
-		body,
-		isBase64Encoded: result.isBase64Encoded,
+		body: body.toString(isBase64Encoded ? "base64" : "utf8"),
+		isBase64Encoded,
 		multiValueHeaders,
 	};
 	debug(response);
 	return response;
 }
 
 export default {
-	convertFrom: convertFromAPIGatewayProxyEvent,
-	convertTo: convertToApiGatewayProxyResult,
+	convertFrom: (event) => convertFromAPIGatewayProxyEvent(event as APIGatewayProxyEvent),
+	convertTo: async () => {
+		const { streamCreator, output } = createBufferedStreamCreator(convertToApiGatewayProxyResult);
+		return { type: "stream" as const, streamCreator, output };
+	},
 	name: "aws-apigw-v1",
-} as Converter;
+} satisfies Converter;

packages/aws/src/overrides/converters/aws-apigw-v2.ts

@@ -5,11 +5,12 @@ import {
 	extractHostFromHeaders,
 	removeUndefinedFromQuery,
 } from "@opennextjs/core/overrides/converters/utils.js";
-import type { InternalEvent, InternalResult } from "@opennextjs/core/types/open-next.js";
+import type { InternalEvent } from "@opennextjs/core/types/open-next.js";
 import type { Converter } from "@opennextjs/core/types/overrides.js";
-import { fromReadableStream } from "@opennextjs/core/utils/stream.js";
 import type { APIGatewayProxyEventV2, APIGatewayProxyResultV2 } from "aws-lambda";
 
+import { createBufferedStreamCreator } from "./response-stream.js";
+
 // Not sure which one is really needed as this is not documented anywhere but server actions redirect are not working without this,
 // it causes a 500 error from cloudfront itself with a 'x-amzErrortype: InternalFailure' header
 const CloudFrontBlacklistedHeaders = [
@@ -68,7 +69,9 @@ function normalizeAPIGatewayProxyEventV2Headers(event: APIGatewayProxyEventV2):
 	return headers;
 }
 
-async function convertFromAPIGatewayProxyEventV2(event: APIGatewayProxyEventV2): Promise<InternalEvent> {
+export async function convertFromAPIGatewayProxyEventV2(
+	event: APIGatewayProxyEventV2
+): Promise<InternalEvent> {
 	const { rawPath, rawQueryString, requestContext } = event;
 	const headers = normalizeAPIGatewayProxyEventV2Headers(event);
 	return {
@@ -92,9 +95,13 @@ async function convertFromAPIGatewayProxyEventV2(event: APIGatewayProxyEventV2):
 	};
 }
 
-async function convertToApiGatewayProxyResultV2(result: InternalResult): Promise<APIGatewayProxyResultV2> {
+function convertToApiGatewayProxyResultV2(
+	prelude: { statusCode: number; cookies: string[]; headers: Record<string, string> },
+	body: Buffer,
+	isBase64Encoded: boolean
+): APIGatewayProxyResultV2 {
 	const headers: Record<string, string> = {};
-	Object.entries(result.headers)
+	Object.entries(prelude.headers)
 		.map(([key, value]) => [key.toLowerCase(), value] as const)
 		.filter(
 			([key]) =>
@@ -110,21 +117,27 @@ async function convertToApiGatewayProxyResultV2(result: InternalResult): Promise
 			headers[key] = Array.isArray(value) ? value.join(", ") : `${value}`;
 		});
 
-	const body = await fromReadableStream(result.body, result.isBase64Encoded);
-
 	const response: APIGatewayProxyResultV2 = {
-		statusCode: result.statusCode,
+		statusCode: prelude.statusCode,
 		headers,
-		cookies: parseSetCookieHeader(result.headers["set-cookie"]),
-		body,
-		isBase64Encoded: result.isBase64Encoded,
+		cookies:
+			prelude.cookies.length > 0
+				? prelude.cookies
+				: prelude.headers["set-cookie"]
+					? parseSetCookieHeader(prelude.headers["set-cookie"])
+					: undefined,
+		body: body.toString(isBase64Encoded ? "base64" : "utf8"),
+		isBase64Encoded,
 	};
 	debug(response);
 	return response;
 }
 
 export default {
-	convertFrom: convertFromAPIGatewayProxyEventV2,
-	convertTo: convertToApiGatewayProxyResultV2,
+	convertFrom: (event) => convertFromAPIGatewayProxyEventV2(event as APIGatewayProxyEventV2),
+	convertTo: async () => {
+		const { streamCreator, output } = createBufferedStreamCreator(convertToApiGatewayProxyResultV2);
+		return { type: "stream" as const, streamCreator, output };
+	},
 	name: "aws-apigw-v2",
-} as Converter;
+} satisfies Converter;

packages/aws/src/overrides/converters/aws-cloudfront.ts

@@ -6,7 +6,6 @@ import { parseSetCookieHeader } from "@opennextjs/core/http/util.js";
 import { extractHostFromHeaders } from "@opennextjs/core/overrides/converters/utils.js";
 import type { InternalEvent, InternalResult, MiddlewareResult } from "@opennextjs/core/types/open-next.js";
 import type { Converter } from "@opennextjs/core/types/overrides.js";
-import { fromReadableStream } from "@opennextjs/core/utils/stream.js";
 import type {
 	CloudFrontCustomOrigin,
 	CloudFrontHeaders,
@@ -15,6 +14,8 @@ import type {
 	CloudFrontRequestResult,
 } from "aws-lambda";
 
+import { createBufferedStreamCreator } from "./response-stream.js";
+
 const cloudfrontBlacklistedHeaders = [
 	// Disallowed headers, see: https://docs.aws.amazon.com/AmazonCloudFront/latest/DeveloperGuide/edge-function-restrictions-all.html#function-restrictions-disallowed-headers
 	"connection",
@@ -129,67 +130,83 @@ function convertToCloudfrontHeaders(headers: Record<string, OutgoingHttpHeader>,
 	return cloudfrontHeaders;
 }
 
-async function convertToCloudFrontRequestResult(
-	result: InternalResult | MiddlewareResult,
+async function convertMiddlewareResult(
+	result: MiddlewareResult,
 	originalRequest: CloudFrontRequestEvent
 ): Promise<CloudFrontRequestResult> {
-	if (result.type === "middleware") {
-		const { method, clientIp, origin } = originalRequest.Records[0].cf.request;
-		const responseHeaders = result.internalEvent.headers;
-
-		// Handle external rewrite
-
-		let customOrigin = origin?.custom as CloudFrontCustomOrigin;
-		let host = responseHeaders.host ?? responseHeaders.Host;
-		if (result.origin) {
-			customOrigin = {
-				...customOrigin,
-				domainName: result.origin.host,
-				port: result.origin.port ?? 443,
-				protocol: result.origin.protocol ?? "https",
-				customHeaders: {},
-			};
-			host = result.origin.host;
-		}
-
-		const response: CloudFrontRequest = {
-			clientIp,
-			method,
-			uri: result.internalEvent.rawPath,
-			querystring: convertToQueryString(result.internalEvent.query).replace("?", ""),
-			headers: convertToCloudfrontHeaders({
-				...responseHeaders,
-				host,
-			}),
-			origin: origin?.custom
-				? {
-						custom: customOrigin,
-					}
-				: origin,
+	const { method, clientIp, origin } = originalRequest.Records[0].cf.request;
+	const responseHeaders = result.internalEvent.headers;
+
+	// Handle external rewrite
+
+	let customOrigin = origin?.custom as CloudFrontCustomOrigin;
+	let host = responseHeaders.host ?? responseHeaders.Host;
+	if (result.origin) {
+		customOrigin = {
+			...customOrigin,
+			domainName: result.origin.host,
+			port: result.origin.port ?? 443,
+			protocol: result.origin.protocol ?? "https",
+			customHeaders: {},
 		};
+		host = result.origin.host;
+	}
 
-		debug("response rewrite", response);
+	const response: CloudFrontRequest = {
+		clientIp,
+		method,
+		uri: result.internalEvent.rawPath,
+		querystring: convertToQueryString(result.internalEvent.query).replace("?", ""),
+		headers: convertToCloudfrontHeaders({
+			...responseHeaders,
+			host,
+		}),
+		origin: origin?.custom
+			? {
+					custom: customOrigin,
+				}
+			: origin,
+	};
 
-		return response;
-	}
+	debug("response rewrite", response);
 
-	const body = await fromReadableStream(result.body, result.isBase64Encoded);
-	const responseHeaders = result.headers;
+	return response;
+}
 
+function convertToCloudFrontRequestResult(
+	prelude: { statusCode: number; cookies: string[]; headers: Record<string, string> },
+	body: Buffer,
+	isBase64Encoded: boolean
+): CloudFrontRequestResult {
+	const responseHeaders = {
+		...prelude.headers,
+		...(prelude.cookies.length > 0 ? { "set-cookie": prelude.cookies } : {}),
+	};
 	const response: CloudFrontRequestResult = {
-		status: result.statusCode.toString(),
+		status: prelude.statusCode.toString(),
 		statusDescription: "OK",
 		headers: convertToCloudfrontHeaders(responseHeaders, true),
-		bodyEncoding: result.isBase64Encoded ? "base64" : "text",
-		body,
+		bodyEncoding: isBase64Encoded ? "base64" : "text",
+		body: body.toString(isBase64Encoded ? "base64" : "utf8"),
 	};
 
 	debug(response);
 	return response;
 }
 
 export default {
-	convertFrom: convertFromCloudFrontRequestEvent,
-	convertTo: convertToCloudFrontRequestResult,
+	convertFrom: (event) => convertFromCloudFrontRequestEvent(event as CloudFrontRequestEvent),
+	convertTo: async (event) => {
+		const { streamCreator, output } = createBufferedStreamCreator(convertToCloudFrontRequestResult);
+		return {
+			type: "stream" as const,
+			streamCreator,
+			output,
+			data: async (result) =>
+				result.type === "middleware"
+					? convertMiddlewareResult(result, event as CloudFrontRequestEvent)
+					: undefined,
+		};
+	},
 	name: "aws-cloudfront",
-} as Converter;
+} satisfies Converter<InternalEvent, InternalResult | MiddlewareResult>;

packages/aws/src/overrides/converters/aws-streaming.ts

@@ -0,0 +1,41 @@
+import type { Writable } from "node:stream";
+
+import type { StreamCreator } from "@opennextjs/core/types/open-next.js";
+import type { Converter } from "@opennextjs/core/types/overrides.js";
+import type { APIGatewayProxyEventV2 } from "aws-lambda";
+
+import { convertFromAPIGatewayProxyEventV2 } from "./aws-apigw-v2.js";
+
+type StreamingContext = {
+	responseStream: Writable & { setContentType(contentType: string): void };
+	writable: Writable;
+	contentEncoding: string;
+};
+
+const converter: Converter = {
+	convertFrom: (event) => convertFromAPIGatewayProxyEventV2(event as APIGatewayProxyEventV2),
+	convertTo: async (_event, context) => {
+		const { responseStream, writable, contentEncoding } = context as StreamingContext;
+		const streamCreator: StreamCreator = {
+			writeHeaders(prelude) {
+				responseStream.setContentType("application/vnd.awslambda.http-integration-response");
+				responseStream.write(
+					JSON.stringify({
+						...prelude,
+						headers: {
+							...prelude.headers,
+							"content-encoding": contentEncoding,
+						},
+					})
+				);
+				responseStream.write(new Uint8Array(8));
+				return writable;
+			},
+		};
+
+		return { type: "stream" as const, streamCreator };
+	},
+	name: "aws-streaming",
+};
+
+export default converter;

packages/aws/src/overrides/converters/response-stream.ts

@@ -0,0 +1,45 @@
+import { Writable } from "node:stream";
+
+import type { StreamCreator } from "@opennextjs/core/types/open-next.js";
+import { isBinaryContentType } from "@opennextjs/core/utils/binary.js";
+
+type Prelude = Parameters<StreamCreator["writeHeaders"]>[0];
+
+export function createBufferedStreamCreator<T>(
+	createOutput: (prelude: Prelude, body: Buffer, isBase64Encoded: boolean) => T
+): { streamCreator: StreamCreator; output: Promise<T> } {
+	const { promise: output, resolve, reject } = Promise.withResolvers<T>();
+	let prelude: Prelude | undefined;
+	const chunks: Buffer[] = [];
+
+	const streamCreator: StreamCreator = {
+		writeHeaders(value) {
+			prelude = value;
+			return new Writable({
+				write(chunk, _encoding, callback) {
+					chunks.push(Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk));
+					callback();
+				},
+				final(callback) {
+					if (!prelude) {
+						const error = new Error("Response stream finished before headers were written");
+						reject(error);
+						callback(error);
+						return;
+					}
+					try {
+						const isBase64Encoded =
+							isBinaryContentType(prelude.headers["content-type"]) || !!prelude.headers["content-encoding"];
+						resolve(createOutput(prelude, Buffer.concat(chunks), isBase64Encoded));
+						callback();
+					} catch (error: unknown) {
+						reject(error);
+						callback(error instanceof Error ? error : new Error(String(error)));
+					}
+				},
+			});
+		},
+	};
+
+	return { streamCreator, output };
+}