feat: support initial-response in RPC based event streams

Package.swift

@@ -555,4 +555,4 @@ let servicesWithIntegrationTests: [String] = [
 servicesWithIntegrationTests.forEach(addIntegrationTestTarget)
 
 // Uncomment this line to enable protocol tests
-// addProtocolTests()
+// addProtocolTests()

Sources/Core/AWSClientRuntime/EventStream/AWSMessageDecoder.swift

@@ -17,6 +17,9 @@ extension AWSEventStream {
         private var decoder: EventStreamMessageDecoder?
         private var messageBuffer: [EventStream.Message] = []
         private var error: Error?
+        private var initialMessage: Data = Data()
+        private var onInitialResponseReceived: ((Data?) -> Void)?
+        private var didProcessInitialMessage = false
 
         private var decodedPayload = Data()
         private var decodededHeaders: [EventStreamHeader] = []
@@ -44,8 +47,20 @@ extension AWSEventStream {
                     self.logger.debug("onComplete")
                     let message = EventStream.Message(headers: self.decodededHeaders.toHeaders(),
                                                       payload: self.decodedPayload)
-                    self.messageBuffer.append(message)
+                    if message.headers.contains(
+                        EventStream.Header(name: ":event-type", value: .string("initial-response"))
+                    ) {
+                        self.initialMessage = message.payload
+                        self.onInitialResponseReceived?(self.initialMessage)
+                        self.didProcessInitialMessage = true
+                    } else {
+                        self.messageBuffer.append(message)
 
+                        if !self.didProcessInitialMessage {
+                            self.onInitialResponseReceived?(nil)  // Signal that initial-response will never come.
+                            self.didProcessInitialMessage = true
+                        }
+                    }
                     // This could be end of the stream, hence reset the state
                     self.decodedPayload = Data()
                     self.decodededHeaders = []
@@ -88,6 +103,38 @@ extension AWSEventStream {
             return message
         }
 
+        // Responsible for waiting on the initial response.
+        // It uses Swift's concurrency model to asynchronously return the data.
+        public func awaitInitialResponse() async -> Data? {
+            // The 'withCheckedContinuation' function is used to bridge asynchronous code
+            // that doesn't use Swift's concurrency model with code that does.
+            return await withCheckedContinuation { continuation in
+                // Here, we attempt to retrieve the initial response.
+                // Once the data is retrieved (or determined to be nil),
+                // the continuation is resumed with the result.
+                retrieveInitialResponse { data in
+                    continuation.resume(returning: data)
+                }
+            }
+        }
+
+        // Attempt to get the initial response.
+        // If the initial message has been processed, it immediately calls the completion handler.
+        // Otherwise, it sets up a callback to be triggered once the initial response is received.
+        private func retrieveInitialResponse(completion: @escaping (Data?) -> Void) {
+            // Check if the initial message has already been processed.
+            if self.didProcessInitialMessage {
+                // If it has been processed, immediately call the completion handler
+                // with the potentially nil or populated 'initialMessage' value.
+                completion(initialMessage)
+            } else {
+                // If the initial message hasn't been processed,
+                // set the 'onInitialResponseReceived' callback to our completion handler,
+                // so it can be called later once the initial response is received.
+                self.onInitialResponseReceived = completion
+            }
+        }
+
         /// Throws an error if one has occurred.
         /// This should be called before any other methods to make sure
         /// that the decoder is in a valid state.

Tests/Core/AWSClientRuntimeTests/EventStream/AWSMessageDecoderTests.swift

@@ -30,6 +30,13 @@ final class AWSMessageDecoderTests: XCTestCase {
         XCTAssertEqual(validMessageNoHeaders, decodedMessage)
     }
 
+    func testDecode_MessageWithInitialResponse() {
+        try! sut.feed(data: validInitialResponseMessageData)
+        let decodedMessage = try! sut.message()
+        // initialResponse message should not be added to the messageBuffer
+        XCTAssertNil(decodedMessage)
+    }
+
     func testEndOfStream_StreamClosed() {
         try! sut.feed(data: validMessageDataNoHeaders[0..<validMessageDataNoHeaders.count-1])
         XCTAssertThrowsError(try sut.endOfStream()) { error in
@@ -40,4 +47,16 @@ final class AWSMessageDecoderTests: XCTestCase {
             XCTAssertEqual("Stream ended before message was complete", message)
         }
     }
+
+    func testAwaitInitialResponse_MessageWithInitialResponse() async {
+        try! sut.feed(data: validInitialResponseMessageData)
+        guard let initialResponse = await sut.awaitInitialResponse() else {
+            XCTFail("Error!")
+            return
+        }
+        XCTAssertEqual(
+            String(data: validInitialResponseMessage.payload, encoding: .utf8),
+            String(data: initialResponse, encoding: .utf8)
+        )
+    }
 }

Tests/Core/AWSClientRuntimeTests/EventStream/AWSMessageEncoderTests.swift

@@ -26,4 +26,9 @@ final class AWSMessageEncoderTests: XCTestCase {
         let encodedMessage = try! sut.encode(message: validMessageNoHeaders)
         XCTAssertEqual(validMessageDataNoHeaders, encodedMessage)
     }
+
+   func testEncode_InitialResponseMessage() {
+       let encodedMessage = try! sut.encode(message: validInitialResponseMessage)
+       XCTAssertEqual(validInitialResponseMessageData, encodedMessage)
+   }
 }

Tests/Core/AWSClientRuntimeTests/EventStream/EventStreamTestData.swift

@@ -92,3 +92,25 @@ let validMessageDataNoHeaders = Data([
     0x68, 0x65, 0x72, 0x20, 0x74, 0x65, 0x73, 0x74, 0x20, 0x70, 0x61, 0x79, 0x6c, 0x6f, 0x61, 0x64,
     0x8d, 0xf8, 0x0e, 0x65,
 ])
+
+let validInitialResponseMessage: EventStream.Message = EventStream.Message(
+    headers: [
+        .init(name: ":event-type", value: .string("initial-response")),
+        .init(name: ":message-type", value: .string("event")),
+        .init(name: ":content-type", value: .string("application/x-amz-json-1.1"))
+    ],
+    payload: "{\"someMetadata\": \"test\"}".data(using: .utf8)!
+)
+
+/// This is the encoded version of the above message validInitialResponseMessage
+let validInitialResponseMessageData = Data([
+    0x00, 0x00, 0x00, 0x88, 0x00, 0x00, 0x00, 0x60, 0xa9, 0x06, 0x45, 0x62, 0x0b, 0x3a, 0x65, 0x76,
+    0x65, 0x6e, 0x74, 0x2d, 0x74, 0x79, 0x70, 0x65, 0x07, 0x00, 0x10, 0x69, 0x6e, 0x69, 0x74, 0x69,
+    0x61, 0x6c, 0x2d, 0x72, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x0d, 0x3a, 0x6d, 0x65, 0x73,
+    0x73, 0x61, 0x67, 0x65, 0x2d, 0x74, 0x79, 0x70, 0x65, 0x07, 0x00, 0x05, 0x65, 0x76, 0x65, 0x6e,
+    0x74, 0x0d, 0x3a, 0x63, 0x6f, 0x6e, 0x74, 0x65, 0x6e, 0x74, 0x2d, 0x74, 0x79, 0x70, 0x65, 0x07,
+    0x00, 0x1a, 0x61, 0x70, 0x70, 0x6c, 0x69, 0x63, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x2f, 0x78, 0x2d,
+    0x61, 0x6d, 0x7a, 0x2d, 0x6a, 0x73, 0x6f, 0x6e, 0x2d, 0x31, 0x2e, 0x31, 0x7b, 0x22, 0x73, 0x6f,
+    0x6d, 0x65, 0x4d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x22, 0x3a, 0x20, 0x22, 0x74, 0x65,
+    0x73, 0x74, 0x22, 0x7d, 0xe6, 0x23, 0x00, 0xf1,
+])