Merge pull request #1332 from mdesousa/input-output-webstream

fix: log web stream responses

Author: willfarrell

packages/input-output-logger/index.js

@@ -1,4 +1,5 @@
 import { Transform } from "node:stream";
+import { TransformStream } from "node:stream/web";
 
 const defaults = {
 	logger: (message) => {
@@ -65,11 +66,19 @@ const inputOutputLoggerMiddleware = (opts = {}) => {
 		omitAndLog("event", request);
 	};
 	const inputOutputLoggerMiddlewareAfter = async (request) => {
+		// Check for Node.js stream
 		if (
 			request.response?._readableState ??
 			request.response?.body?._readableState
 		) {
 			passThrough(request, omitAndLog);
+		}
+		// Check for Web stream
+		else if (
+			request.response instanceof ReadableStream ||
+			request.response?.body instanceof ReadableStream
+		) {
+			passThroughWebStream(request, omitAndLog);
 		} else {
 			omitAndLog("response", request);
 		}
@@ -162,4 +171,39 @@ const passThrough = (request, omitAndLog) => {
 	}
 };
 
+// Handler for Web Streams API
+const passThroughWebStream = (request, omitAndLog) => {
+	const hasBody = request.response?.body;
+	let body = "";
+
+	const transformer = new TransformStream({
+		transform(chunk, controller) {
+			// For web streams, chunks could be various types
+			const textChunk =
+				typeof chunk === "string"
+					? chunk
+					: chunk instanceof Uint8Array
+						? new TextDecoder().decode(chunk)
+						: String(chunk);
+			body += textChunk;
+			controller.enqueue(chunk);
+		},
+		flush(controller) {
+			if (hasBody) {
+				omitAndLog("response", { response: { ...request.response, body } });
+			} else {
+				omitAndLog("response", { response: body });
+			}
+		},
+	});
+
+	if (hasBody) {
+		// Handle response with body property that's a ReadableStream
+		request.response.body = request.response.body.pipeThrough(transformer);
+	} else {
+		// Handle response that's directly a ReadableStream
+		request.response = request.response.pipeThrough(transformer);
+	}
+};
+
 export default inputOutputLoggerMiddleware;

packages/input-output-logger/index.test.js

@@ -116,6 +116,96 @@ test("It should log with streamifyResponse:true using body ReadableStream", asyn
 	]);
 });
 
+test("It should log with Web Streams API using ReadableStream", async (t) => {
+	const input = "x".repeat(1024 * 1024);
+	const logger = t.mock.fn();
+	const handler = middy(
+		async (event, context, { signal }) => {
+			// Create a Web ReadableStream
+			const stream = new ReadableStream({
+				start(controller) {
+					controller.enqueue(input);
+					controller.close();
+				},
+			});
+			return stream;
+		},
+		{
+			streamifyResponse: true,
+		},
+	).use(
+		inputOutputLogger({
+			logger,
+		}),
+	);
+
+	const event = {};
+	let chunkResponse = "";
+	const responseStream = createWritableStream((chunk) => {
+		chunkResponse += chunk;
+	});
+	const response = await handler(event, responseStream, context);
+	equal(response, undefined);
+	equal(chunkResponse, input);
+	deepEqual(logger.mock.calls[0].arguments, [{ event: {} }]);
+	deepEqual(logger.mock.calls[1].arguments, [
+		{
+			response: input,
+		},
+	]);
+});
+
+test("It should log with Web Streams API using body ReadableStream", async (t) => {
+	const input = "x".repeat(1024 * 1024);
+	const logger = t.mock.fn();
+	const handler = middy(
+		async (event, context, { signal }) => {
+			// Create a Web ReadableStream in the body
+			const stream = new ReadableStream({
+				start(controller) {
+					controller.enqueue(input);
+					controller.close();
+				},
+			});
+			return {
+				statusCode: 200,
+				headers: {
+					"Content-Type": "plain/text",
+				},
+				body: stream,
+			};
+		},
+		{
+			streamifyResponse: true,
+		},
+	).use(
+		inputOutputLogger({
+			logger,
+		}),
+	);
+
+	const event = {};
+	let chunkResponse = "";
+	const responseStream = createWritableStream((chunk) => {
+		chunkResponse += chunk;
+	});
+	const response = await handler(event, responseStream, context);
+	equal(response, undefined);
+	equal(chunkResponse, input);
+	deepEqual(logger.mock.calls[0].arguments, [{ event: {} }]);
+	deepEqual(logger.mock.calls[1].arguments, [
+		{
+			response: {
+				statusCode: 200,
+				headers: {
+					"Content-Type": "plain/text",
+				},
+				body: input,
+			},
+		},
+	]);
+});
+
 test("It should throw error when invalid logger", async (t) => {
 	const logger = false;