add support for streaming lambda function + custom status code and HTTP headers

Author: sebsto

.gitignore

@@ -13,4 +13,5 @@ Package.resolved
 Makefile
 .devcontainer
 .amazonq
-.kiro
+.kiro
+nodejs

Examples/Streaming/Sources/main.swift

@@ -31,21 +31,19 @@ struct SendNumbersWithPause: StreamingLambdaHandler {
         // Send HTTP status code and headers before streaming the response body
         try await responseWriter.writeStatusAndHeaders(
             StreamingLambdaStatusAndHeadersResponse(
-                statusCode: 200,
+                statusCode: 418,  // I'm a tea pot
                 headers: [
                     "Content-Type": "text/plain",
                     "x-my-custom-header": "streaming-example",
-                ],
-                multiValueHeaders: [
-                    "Set-Cookie": ["session=abc123", "theme=dark"]
                 ]
             )
         )
 
         // Stream numbers with pauses to demonstrate streaming functionality
-        for i in 1...10 {
+        for i in 1...3 {
             // Send partial data
             try await responseWriter.write(ByteBuffer(string: "Number: \(i)\n"))
+
             // Perform some long asynchronous work to simulate processing
             try await Task.sleep(for: .milliseconds(1000))
         }
@@ -58,64 +56,5 @@ struct SendNumbersWithPause: StreamingLambdaHandler {
     }
 }
 
-// Example of a more complex streaming handler that demonstrates different response scenarios
-struct ConditionalStreamingHandler: StreamingLambdaHandler {
-    func handle(
-        _ event: ByteBuffer,
-        responseWriter: some LambdaResponseStreamWriter,
-        context: LambdaContext
-    ) async throws {
-
-        // Parse the event to determine response type
-        let eventString = String(buffer: event)
-        let shouldError = eventString.contains("error")
-
-        if shouldError {
-            // Send error response with appropriate status code
-            try await responseWriter.writeStatusAndHeaders(
-                StreamingLambdaStatusAndHeadersResponse(
-                    statusCode: 400,
-                    headers: [
-                        "Content-Type": "application/json",
-                        "x-error-type": "client-error",
-                    ]
-                )
-            )
-
-            try await responseWriter.writeAndFinish(
-                ByteBuffer(string: #"{"error": "Bad request", "message": "Error requested in input"}"#)
-            )
-        } else {
-            // Send successful response with streaming data
-            try await responseWriter.writeStatusAndHeaders(
-                StreamingLambdaStatusAndHeadersResponse(
-                    statusCode: 200,
-                    headers: [
-                        "Content-Type": "application/json",
-                        "Cache-Control": "no-cache",
-                    ]
-                )
-            )
-
-            // Stream JSON array elements
-            try await responseWriter.write(ByteBuffer(string: "["))
-
-            for i in 1...5 {
-                if i > 1 {
-                    try await responseWriter.write(ByteBuffer(string: ","))
-                }
-                try await responseWriter.write(
-                    ByteBuffer(string: #"{"id": \#(i), "timestamp": "\#(Date().timeIntervalSince1970)"}"#)
-                )
-                try await Task.sleep(for: .milliseconds(500))
-            }
-
-            try await responseWriter.write(ByteBuffer(string: "]"))
-            try await responseWriter.finish()
-        }
-    }
-}
-
-// Use the simple example by default
 let runtime = LambdaRuntime(handler: SendNumbersWithPause())
 try await runtime.run()

Examples/Streaming/template.yaml

@@ -17,6 +17,9 @@ Resources:
       FunctionUrlConfig:
         AuthType: AWS_IAM
         InvokeMode: RESPONSE_STREAM
+      Environment:
+        Variables:
+          LOG_LEVEL: trace
 
 Outputs:
   # print Lambda function URL

Sources/AWSLambdaRuntime/LambdaHandlers.swift

@@ -57,6 +57,15 @@ public protocol LambdaResponseStreamWriter {
     /// Write a response part into the stream and then end the stream as well as the underlying HTTP response.
     /// - Parameter buffer: The buffer to write.
     func writeAndFinish(_ buffer: ByteBuffer) async throws
+
+    /// Write a response part into the stream.
+    // In the context of streaming Lambda, this is used to allow the user
+    // to send custom headers or statusCode.
+    /// - Note: user should use the writeStatusAndHeaders(:StreamingLambdaStatusAndHeadersResponse)
+    //          function to write the status code and headers
+    /// - Parameter buffer: The buffer corresponding to the status code and headers to write.
+    func writeCustomHeader(_ buffer: NIOCore.ByteBuffer) async throws
+
 }
 
 /// This handler protocol is intended to serve the most common use-cases.

Sources/AWSLambdaRuntime/LambdaResponseStreamWriter+Headers.swift

@@ -44,7 +44,7 @@ public struct StreamingLambdaStatusAndHeadersResponse: Codable, Sendable {
     public init(
         statusCode: Int,
         headers: [String: String]? = nil,
-        multiValueHeaders: [String: [String]]? = nil
+        multiValueHeaders: [String: [String]]? = nil,
     ) {
         self.statusCode = statusCode
         self.headers = headers
@@ -67,17 +67,15 @@ extension LambdaResponseStreamWriter {
         encoder: Encoder
     ) async throws where Encoder.Output == StreamingLambdaStatusAndHeadersResponse {
 
-        // Convert Data to ByteBuffer
+        // Convert JSON headers to an array of bytes in a ByteBuffer
         var buffer = ByteBuffer()
         try encoder.encode(response, into: &buffer)
 
-        // Write the JSON data
-        try await write(buffer)
-
         // Write eight null bytes as separator
-        var separatorBuffer = ByteBuffer()
-        separatorBuffer.writeBytes([0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00])
-        try await write(separatorBuffer)
+        buffer.writeBytes([0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00])
+
+        // Write the JSON data and the separator
+        try await writeCustomHeader(buffer)
     }
 }
 
@@ -95,6 +93,7 @@ extension LambdaResponseStreamWriter {
         _ response: StreamingLambdaStatusAndHeadersResponse,
         encoder: JSONEncoder = JSONEncoder()
     ) async throws {
+        encoder.outputFormatting = .withoutEscapingSlashes
         try await self.writeStatusAndHeaders(response, encoder: LambdaJSONOutputEncoder(encoder))
     }
 }

Sources/AWSLambdaRuntime/LambdaRuntimeClient.swift

@@ -16,9 +16,13 @@ import Logging
 import NIOCore
 import NIOHTTP1
 import NIOPosix
+import Synchronization
 
 @usableFromInline
 final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {
+    @usableFromInline
+    var _hasStreamingCustomHeaders = false
+
     @usableFromInline
     nonisolated let unownedExecutor: UnownedSerialExecutor
 
@@ -42,6 +46,11 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {
             self.runtimeClient = runtimeClient
         }
 
+        @usableFromInline
+        func writeCustomHeader(_ buffer: NIOCore.ByteBuffer) async throws {
+            try await self.runtimeClient.writeCustomHeader(buffer)
+        }
+
         @usableFromInline
         func write(_ buffer: NIOCore.ByteBuffer) async throws {
             try await self.runtimeClient.write(buffer)
@@ -188,6 +197,10 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {
         }
     }
 
+    private func writeCustomHeader(_ buffer: NIOCore.ByteBuffer) async throws {
+        _hasStreamingCustomHeaders = true
+        try await self.write(buffer)
+    }
     private func write(_ buffer: NIOCore.ByteBuffer) async throws {
         switch self.lambdaState {
         case .idle, .sentResponse:
@@ -210,6 +223,7 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {
     }
 
     private func writeAndFinish(_ buffer: NIOCore.ByteBuffer?) async throws {
+        _hasStreamingCustomHeaders = false
         switch self.lambdaState {
         case .idle, .sentResponse:
             throw LambdaRuntimeError(code: .finishAfterFinishHasBeenSent)
@@ -330,7 +344,11 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {
                         NIOHTTPClientResponseAggregator(maxContentLength: 6 * 1024 * 1024)
                     )
                     try channel.pipeline.syncOperations.addHandler(
-                        LambdaChannelHandler(delegate: self, logger: self.logger, configuration: self.configuration)
+                        LambdaChannelHandler(
+                            delegate: self,
+                            logger: self.logger,
+                            configuration: self.configuration
+                        )
                     )
                     return channel.eventLoop.makeSucceededFuture(())
                 } catch {
@@ -425,13 +443,17 @@ extension LambdaRuntimeClient: LambdaChannelHandlerDelegate {
 
             }
         }
+    }
 
+    func hasStreamingCustomHeaders(isolation: isolated (any Actor)? = #isolation) async -> Bool {
+        await self._hasStreamingCustomHeaders
     }
 }
 
 private protocol LambdaChannelHandlerDelegate {
     func connectionWillClose(channel: any Channel)
     func connectionErrorHappened(_ error: any Error, channel: any Channel)
+    func hasStreamingCustomHeaders(isolation: isolated (any Actor)?) async -> Bool
 }
 
 private final class LambdaChannelHandler<Delegate: LambdaChannelHandlerDelegate> {
@@ -467,10 +489,16 @@ private final class LambdaChannelHandler<Delegate: LambdaChannelHandlerDelegate>
     let defaultHeaders: HTTPHeaders
     /// These headers must be sent along an invocation or initialization error report
     let errorHeaders: HTTPHeaders
-    /// These headers must be sent when streaming a response
+    /// These headers must be sent when streaming a large response
+    let largeResponseHeaders: HTTPHeaders
+    /// These headers must be sent when the handler streams its response
     let streamingHeaders: HTTPHeaders
 
-    init(delegate: Delegate, logger: Logger, configuration: LambdaRuntimeClient.Configuration) {
+    init(
+        delegate: Delegate,
+        logger: Logger,
+        configuration: LambdaRuntimeClient.Configuration
+    ) {
         self.delegate = delegate
         self.logger = logger
         self.configuration = configuration
@@ -483,11 +511,23 @@ private final class LambdaChannelHandler<Delegate: LambdaChannelHandlerDelegate>
             "user-agent": .userAgent,
             "lambda-runtime-function-error-type": "Unhandled",
         ]
-        self.streamingHeaders = [
+        self.largeResponseHeaders = [
             "host": "\(self.configuration.ip):\(self.configuration.port)",
             "user-agent": .userAgent,
             "transfer-encoding": "chunked",
         ]
+        // https://docs.aws.amazon.com/lambda/latest/dg/runtimes-custom.html#runtimes-custom-response-streaming
+        // These are the headers returned by the Runtime to the Lambda Data plane.
+        // These are not the headers the Lambda Data plane sends to the caller of the Lambda function
+        // The developer of the function can set the caller's headers in the handler code.
+        self.streamingHeaders = [
+            "host": "\(self.configuration.ip):\(self.configuration.port)",
+            "user-agent": .userAgent,
+            "Lambda-Runtime-Function-Response-Mode": "streaming",
+            // these are not used by this runtime client at the moment
+            // FIXME: the eror handling should inject these headers in the streamed response to report mid-stream errors
+            "Trailer": "Lambda-Runtime-Function-Error-Type, Lambda-Runtime-Function-Error-Body",
+        ]
     }
 
     func nextInvocation(isolation: isolated (any Actor)? = #isolation) async throws -> Invocation {
@@ -625,11 +665,16 @@ private final class LambdaChannelHandler<Delegate: LambdaChannelHandlerDelegate>
             // TODO: This feels super expensive. We should be able to make this cheaper. requestIDs are fixed length
             let url = Consts.invocationURLPrefix + "/" + requestID + Consts.postResponseURLSuffix
 
+            var headers = self.streamingHeaders
+            if await self.delegate.hasStreamingCustomHeaders(isolation: #isolation) {
+                // this headers is required by Function URL when the user sends custom status code or headers
+                headers.add(name: "Content-Type", value: "application/vnd.awslambda.http-integration-response")
+            }
             let httpRequest = HTTPRequestHead(
                 version: .http1_1,
                 method: .POST,
                 uri: url,
-                headers: self.streamingHeaders
+                headers: headers
             )
 
             context.write(self.wrapOutboundOut(.head(httpRequest)), promise: nil)
@@ -652,17 +697,13 @@ private final class LambdaChannelHandler<Delegate: LambdaChannelHandlerDelegate>
 
             // If we have less than 6MB, we don't want to use the streaming API. If we have more
             // than 6MB we must use the streaming mode.
-            let headers: HTTPHeaders =
-                if byteBuffer?.readableBytes ?? 0 < 6_000_000 {
-                    [
-                        "host": "\(self.configuration.ip):\(self.configuration.port)",
-                        "user-agent": .userAgent,
-                        "content-length": "\(byteBuffer?.readableBytes ?? 0)",
-                    ]
-                } else {
-                    self.streamingHeaders
-                }
-
+            var headers: HTTPHeaders!
+            if byteBuffer?.readableBytes ?? 0 < 6_000_000 {
+                headers = self.defaultHeaders
+                headers.add(name: "content-length", value: "\(byteBuffer?.readableBytes ?? 0)")
+            } else {
+                headers = self.largeResponseHeaders
+            }
             let httpRequest = HTTPRequestHead(
                 version: .http1_1,
                 method: .POST,

Sources/AWSLambdaRuntime/LambdaRuntimeClientProtocol.swift

@@ -20,6 +20,7 @@ package protocol LambdaRuntimeClientResponseStreamWriter: LambdaResponseStreamWr
     func finish() async throws
     func writeAndFinish(_ buffer: ByteBuffer) async throws
     func reportError(_ error: any Error) async throws
+    func writeCustomHeader(_ buffer: NIOCore.ByteBuffer) async throws
 }
 
 @usableFromInline

Tests/AWSLambdaRuntimeTests/Lambda+CodableTests.swift

@@ -96,5 +96,11 @@ struct JSONTests {
         func finish() async throws {
             fatalError("Unexpected call")
         }
+
+        func writeCustomHeader(_ buffer: NIOCore.ByteBuffer) async throws {
+            // This is a mock, so we don't actually write custom headers.
+            // In a real implementation, this would handle writing custom headers.
+            fatalError("Unexpected call to writeCustomHeader")
+        }
     }
 }

Tests/AWSLambdaRuntimeTests/LambdaResponseStreamWriter+HeadersTests.swift

@@ -587,6 +587,12 @@ final class MockLambdaResponseStreamWriter: LambdaResponseStreamWriter {
         writtenBuffers.append(buffer)
         isFinished = true
     }
+
+    func writeCustomHeader(_ buffer: NIOCore.ByteBuffer) async throws {
+        // This is a mock, so we don't actually write custom headers.
+        // In a real implementation, this would handle writing custom headers.
+        fatalError("Unexpected call to writeCustomHeader")
+    }
 }
 
 // MARK: - Error Handling Mock Implementations
@@ -620,6 +626,12 @@ final class FailingMockLambdaResponseStreamWriter: LambdaResponseStreamWriter {
         try await write(buffer)
         try await finish()
     }
+
+    func writeCustomHeader(_ buffer: NIOCore.ByteBuffer) async throws {
+        // This is a mock, so we don't actually write custom headers.
+        // In a real implementation, this would handle writing custom headers.
+        fatalError("Unexpected call to writeCustomHeader")
+    }
 }
 
 // MARK: - Test Error Types
@@ -716,6 +728,12 @@ final class TrackingLambdaResponseStreamWriter: LambdaResponseStreamWriter {
         writtenBuffers.append(buffer)
         isFinished = true
     }
+
+    func writeCustomHeader(_ buffer: NIOCore.ByteBuffer) async throws {
+        // This is a mock, so we don't actually write custom headers.
+        // In a real implementation, this would handle writing custom headers.
+        fatalError("Unexpected call to writeCustomHeader")
+    }
 }
 
 /// Mock implementation with custom behavior for integration testing
@@ -739,4 +757,10 @@ final class CustomBehaviorLambdaResponseStreamWriter: LambdaResponseStreamWriter
         writtenBuffers.append(buffer)
         isFinished = true
     }
+
+    func writeCustomHeader(_ buffer: NIOCore.ByteBuffer) async throws {
+        // This is a mock, so we don't actually write custom headers.
+        // In a real implementation, this would handle writing custom headers.
+        fatalError("Unexpected call to writeCustomHeader")
+    }
 }