W 21683654 improve streaming error handling (#345)

* fix error handling in lambda adapter

* update tests to check 502 prevention behaviour when streaming

* add changeset for fix

Author: noahadams

.changeset/hungry-otters-unite.md

@@ -0,0 +1,5 @@
+---
+'@salesforce/mrt-utilities': patch
+---
+
+Improved error handling in create-lambda-adapter.ts used for streaming support when deployed to Managed Runtime.

packages/mrt-utilities/src/streaming/create-lambda-adapter.ts

@@ -115,7 +115,7 @@ export function createStreamingLambdaAdapter(
 
       if (isStreamOpen && typeof responseStream.write === 'function') {
         const errorMessage = error instanceof Error ? error.message : String(error);
-        responseStream.write(`HTTP/1.1 500 Internal Server Error\r\n\r\nInternal Server Error: ${errorMessage}`);
+        responseStream.write(`Internal Server Error: ${errorMessage}`);
       } else {
         console.error('[error handler] Cannot write error - stream is closed');
       }
@@ -163,34 +163,20 @@ async function streamResponse(
       }
     };
 
+    // See events defined for hhtp.serverResponse interface here:
+    // https://nodejs.org/docs/latest-v24.x/api/http.html#class-httpserverresponse
+
     // Handle response finish
-    expressResponse.once('finish', () => {
-      resolveOnce();
-    });
+    expressResponse.once('finish', resolveOnce);
+    // handle response close, either after a 'finish' or if underlying
+    // connection is closed early
+    expressResponse.once('close', resolveOnce);
 
     // Handle response errors
-    expressResponse.once('error', (err: Error) => {
-      rejectOnce(err);
-    });
+    expressResponse.once('error', rejectOnce);
 
     try {
-      app(expressRequest, expressResponse, (err) => {
-        if (err) {
-          console.error('Express app error:', err);
-          rejectOnce(err);
-        } else {
-          // If response has finished, resolveOnce will be called by the finish event
-          // Otherwise, resolve after a short delay to allow async operations
-          if (expressResponse.finished) {
-            resolveOnce();
-          } else {
-            // Wait a bit for the response to finish
-            setTimeout(() => {
-              resolveOnce();
-            }, 10);
-          }
-        }
-      });
+      app(expressRequest, expressResponse);
     } catch (error) {
       console.error('Error in streamResponse:', error);
       rejectOnce(error as Error);
@@ -530,6 +516,7 @@ export function createExpressResponse(
 
     if (!isStreamOpen()) {
       console.error('Cannot initialize response - stream is closed');
+      emitCloseOnce(response);
       return;
     }
 
@@ -557,7 +544,7 @@ export function createExpressResponse(
     // Create HttpResponseStream with metadata
     // This writes the HTTP status and headers to the stream
     const metadata: StreamMetadata = {
-      statusCode,
+      statusCode: response.statusCode || statusCode,
       headers,
     };
 
@@ -570,6 +557,15 @@ export function createExpressResponse(
 
     httpResponseStream = awslambda.HttpResponseStream.from(responseStream, metadata);
 
+    // Bind 'finish' and 'close' handlers to emit the same events on the response object
+    httpResponseStream.on('finish', () => {
+      response.finished = true;
+      response.emit('finish');
+    });
+    httpResponseStream.on('close', () => {
+      response.emit('close');
+    });
+
     // Set up compression stream if compression is enabled
     // The compression stream pipes to httpResponseStream, which pipes to responseStream
     // 'identity' means no encoding, so we should not initialize compression for it
@@ -670,7 +666,7 @@ export function createExpressResponse(
       reasonPhrase = undefined;
     }
 
-    statusCode = code || statusCode;
+    statusCode = code || this.statusCode || statusCode;
     this.statusCode = statusCode;
 
     // Set statusMessage if provided
@@ -699,6 +695,7 @@ export function createExpressResponse(
   res.write = function (chunk: string | Buffer | Uint8Array): boolean {
     if (!isStreamOpen()) {
       console.error(`Cannot write - stream is closed`);
+      emitCloseOnce(this);
       return false;
     }
 
@@ -729,6 +726,15 @@ export function createExpressResponse(
     }
   };
 
+  let emittedClose = false;
+  const emitCloseOnce = (response: ExpressResponse) => {
+    if (emittedClose) {
+      return;
+    }
+    response.emit('close');
+    emittedClose = true;
+  };
+
   /**
    * Ends the appropriate stream(s) and emits the finish event
    * If compression is enabled, ends the compression stream which will automatically
@@ -737,47 +743,29 @@ export function createExpressResponse(
    * @param response - The Express response object to emit finish event on
    */
   const endStream = (response: ExpressResponse): void => {
-    if (shouldCompress && compressionStream) {
-      try {
-        // Flush compression stream to ensure all buffered data is written
-        _flush();
-        // End compression stream - this will automatically end httpResponseStream
-        // due to the pipe with { end: true } option
-        compressionStream.end(() => {
-          response.finished = true;
-          response.emit('finish');
-        });
-      } catch (error) {
-        console.error(`Error ending compression stream:`, error);
-        // Still emit finish even if there was an error
-        response.finished = true;
-        response.emit('finish');
-      }
-    } else if (httpResponseStream && httpResponseStream.writable) {
-      // No compression, end httpResponseStream directly
+    const stream = (shouldCompress && compressionStream) || (httpResponseStream?.writable && httpResponseStream);
+    if (stream) {
       try {
         _flush();
-        httpResponseStream.end(() => {
-          response.finished = true;
-          response.emit('finish');
-        });
+        stream.end();
       } catch (error) {
-        console.error(`Error ending httpResponseStream:`, error);
+        console.error(`Error ending stream:`, error);
         response.finished = true;
-        response.emit('finish');
+        emitCloseOnce(response);
       }
     } else {
       console.error(`Cannot call end() - stream is closed`);
       // Still emit finish to prevent hanging
       response.finished = true;
-      response.emit('finish');
+      emitCloseOnce(response);
     }
   };
 
   // @ts-expect-error - Type signature doesn't match ServerResponse.end exactly, but our implementation is compatible
   res.end = function (chunk?: string | Buffer | Uint8Array) {
     if (!isStreamOpen()) {
       console.error(`Cannot end - stream is already closed`);
+      emitCloseOnce(this);
       return this;
     }
 
@@ -808,9 +796,8 @@ export function createExpressResponse(
   // Add Express-specific methods that aren't in ServerResponse
   res.status = function (code: number, message?: string) {
     this.statusCode = code;
-    statusCode = code;
     if (message !== undefined) {
-      statusMessage = message;
+      this.statusMessage = message;
     }
     return this;
   };
@@ -855,6 +842,7 @@ export function createExpressResponse(
     if (!responseStarted) {
       if (!isStreamOpen()) {
         console.error('[res.flushHeaders] Cannot flush headers - stream is closed');
+        emitCloseOnce(this);
         return this;
       }
 
@@ -901,6 +889,7 @@ export function createExpressResponse(
   res.flush = function () {
     if (!isStreamOpen()) {
       console.error(`Cannot flush - stream is closed`);
+      emitCloseOnce(this);
       return this;
     }
 

packages/mrt-utilities/test/streaming/create-lambda-adapter.test.ts

@@ -17,12 +17,22 @@ import {
   createExpressResponse,
 } from '@salesforce/mrt-utilities/streaming';
 
+type WritableWithChunksAndMetadata = Writable &
+  EventEmitter & {
+    chunks: Buffer[];
+    metadata: {statusCode: number; headers: Record<string, any>; cookies?: string[]};
+    getWrittenData: (encoding?: BufferEncoding) => Buffer | string;
+  };
+
 // Mock awslambda global
 const mockHttpResponseStream = {
   from: sinon
     .stub()
     .callsFake(
-      (stream: Writable, _metadata: {statusCode: number; headers: Record<string, any>; cookies?: string[]}) => {
+      (
+        stream: WritableWithChunksAndMetadata,
+        _metadata: {statusCode: number; headers: Record<string, any>; cookies?: string[]}) => {
+        stream.metadata = _metadata;
         return stream;
       },
     ),
@@ -35,7 +45,7 @@ const mockHttpResponseStream = {
 };
 
 // Mock stream type with Sinon stubs so assertions (e.g. .called, .calledWith) type-check
-type MockWritable = (Writable & EventEmitter) & {
+type MockWritable = WritableWithChunksAndMetadata & {
   write: sinon.SinonStub;
   end: sinon.SinonStub;
   destroy: sinon.SinonStub;
@@ -53,6 +63,8 @@ function createMockWritable(): MockWritable {
   stream.writableEnded = false;
   stream.writableFinished = false;
   stream.destroyed = false;
+  stream.chunks = chunks;
+  stream.metadata = undefined;
 
   stream.write = sinon.stub().callsFake((chunk: any) => {
     if (destroyed || ended) return false;
@@ -84,6 +96,11 @@ function createMockWritable(): MockWritable {
     // Mock flush method
   });
 
+  stream.getWrittenData = function (encoding?: BufferEncoding) {
+    const data = Buffer.concat(this.chunks);
+    return encoding ? data.toString(encoding) : data;
+};
+
   return stream as MockWritable;
 }
 
@@ -241,7 +258,6 @@ describe('create-lambda-adapter', () => {
         mockResponseStream = createMockWritable();
         mockApp = express();
         mockHttpResponseStream.from.resetHistory();
-        mockHttpResponseStream.from.callsFake((stream) => stream);
       });
 
       afterEach(() => {
@@ -283,7 +299,7 @@ describe('create-lambda-adapter', () => {
 
       await handler(event, context);
 
-      expect(mockResponseStream.write.firstCall.args[0]).to.include('500 Internal Server Error');
+      expect(mockResponseStream.write.firstCall.args[0]).to.include('Error');
       expect(mockResponseStream.end.called).to.be.true;
     });
 
@@ -298,7 +314,7 @@ describe('create-lambda-adapter', () => {
 
       await handler(event, context);
 
-      expect(mockResponseStream.write.firstCall.args[0]).to.include('500 Internal Server Error');
+      expect(mockResponseStream.write.firstCall.args[0]).to.include('Error');
       expect(mockResponseStream.end.called).to.be.true;
     });
 
@@ -321,22 +337,80 @@ describe('create-lambda-adapter', () => {
       expect(closedStream.write.called).to.be.false;
     });
 
-    it('should handle stream without write method', async () => {
-      mockApp.get('/test', () => {
-        throw new Error('Test error');
+    it('should return status code 400 response for requests to streaming route with invalid path', async () => {
+      // We need a catch-all route to force the router to try to decode the path
+      const dummyCatchAllRoute = (_: any, res: any) => {
+        res.status(200).send('dummy catch-all route response');
+      };
+      try {
+        //express 4 style catch-all route, throws an error when installed 
+        // in express 5, see https://github.com/pillarjs/path-to-regexp#errors
+        mockApp.get('/*', dummyCatchAllRoute);
+      } catch (error) {
+        //express 5 style catch-all route
+        mockApp.get('/{*splat}', dummyCatchAllRoute);
+      }
+
+
+      const handler = createStreamingLambdaAdapter(mockApp, mockResponseStream);
+      const event = createMockEvent({path: '/%80'});
+      const context = createMockContext();
+
+      await handler(event, context);
+
+      expect(mockResponseStream.metadata.statusCode).to.equal(400);
+      expect(mockResponseStream.write.called).to.be.true;
+      expect(mockResponseStream.end.called).to.be.true;
+
+      const responseData = mockResponseStream.getWrittenData('utf-8');
+
+      // The response body returned by finalhandler depends on the NODE_ENV environment variable, when it is set to "production", a brief error page with the message "Bad Request" is returned, otherwise it returns a more detailed error page with a stack trace for a URIError exception
+      expect(responseData).to.contain('<title>Error</title>');
+      expect(responseData).to.match(/URIError|Bad Request/);
+    });
+
+    it('should return status code 404 response for requests to streaming route that explicitly raises to return a 404', async () => {
+      // A 404 implemented by throwing
+      mockApp.get('/intentional404', () => {
+        type ErrorWithStatus = Error & {status: number};
+        const err = new Error('Not Found') as ErrorWithStatus;
+        err.message = 'Not Found';
+        err.status = 404;
+        throw err;
       });
 
-      const streamWithoutWrite = createMockWritable();
-      delete (streamWithoutWrite as any).write;
+      const handler = createStreamingLambdaAdapter(mockApp, mockResponseStream);
+      const event = createMockEvent({path: '/intentional404'});
+      const context = createMockContext();
+
+      await handler(event, context);
 
-      const handler = createStreamingLambdaAdapter(mockApp, streamWithoutWrite);
-      const event = createMockEvent({path: '/test'});
+      //expect(mockResponseStream.write).toHaveBeenCalled();
+      expect(mockResponseStream.end.called).to.be.true;
+      expect(mockResponseStream.metadata.statusCode).to.equal(404);
+
+      expect(mockResponseStream.getWrittenData('utf-8')).to.contain('Not Found');
+    });
+
+    it('should stream status code 200 response for requests to an finalhandler style route', async () => {
+      mockApp.get('/simpleroute', (req, res) => {
+        // A hanlder that uses the response object in the style of https://github.com/pillarjs/finalhandler/blob/v2.1.0/index.js#L245-L280 which is used by express to handle the final response in case of an error
+        res.statusCode = 200;
+        res.statusMessage = 'OK';
+        const body = 'hello world';
+        res.setHeader('Content-Type', 'text/plain; charset=utf-8');
+        res.setHeader('Content-Length', Buffer.byteLength(body, 'utf8'));
+        res.end(body);
+      });
+
+      const handler = createStreamingLambdaAdapter(mockApp, mockResponseStream);
+      const event = createMockEvent({path: '/simpleroute'});
       const context = createMockContext();
 
       await handler(event, context);
 
-      // Should not throw
-      expect(streamWithoutWrite.end.called).to.be.true;
+      expect(mockResponseStream.metadata.statusCode).to.equal(200);
+      expect(mockResponseStream.getWrittenData('utf-8')).to.equal('hello world');
     });
 
     it('should handle stream without end method in finally', async () => {