feat!: Remove onMessage callback

This is not necessary anymore, since you can just map the incoming AsyncStream

Author: NeedleInAJayStack

Sources/GraphQLWS/Client.swift

@@ -12,7 +12,6 @@ public actor Client<InitPayload: Equatable & Codable> {
     let onData: (DataResponse, Client) async throws -> Void
     let onError: (ErrorResponse, Client) async throws -> Void
     let onComplete: (CompleteResponse, Client) async throws -> Void
-    let onMessage: (String, Client) async throws -> Void
 
     let encoder = GraphQLJSONEncoder()
     let decoder = JSONDecoder()
@@ -27,16 +26,14 @@ public actor Client<InitPayload: Equatable & Codable> {
     ///   - onData: The callback run on receipt of a `data` message
     ///   - onError: The callback run on receipt of an `error` message
     ///   - onComplete: The callback run on receipt of a `complete` message
-    ///   - onMessage: The callback run on receipt of any message
     public init(
         messenger: Messenger,
         onConnectionError: @escaping (ConnectionErrorResponse, Client) async throws -> Void = { _, _ in },
         onConnectionAck: @escaping (ConnectionAckResponse, Client) async throws -> Void = { _, _ in },
         onConnectionKeepAlive: @escaping (ConnectionKeepAliveResponse, Client) async throws -> Void = { _, _ in },
         onData: @escaping (DataResponse, Client) async throws -> Void = { _, _ in },
         onError: @escaping (ErrorResponse, Client) async throws -> Void = { _, _ in },
-        onComplete: @escaping (CompleteResponse, Client) async throws -> Void = { _, _ in },
-        onMessage: @escaping (String, Client) async throws -> Void = { _, _ in }
+        onComplete: @escaping (CompleteResponse, Client) async throws -> Void = { _, _ in }
     ) {
         self.messenger = messenger
         self.onConnectionError = onConnectionError
@@ -45,15 +42,12 @@ public actor Client<InitPayload: Equatable & Codable> {
         self.onData = onData
         self.onError = onError
         self.onComplete = onComplete
-        self.onMessage = onMessage
     }
 
     /// Listen and react to the provided async sequence of server messages. This function will block until the stream is completed.
     /// - Parameter incoming: The server message sequence that the client should react to.
     public func listen<A: AsyncSequence & Sendable>(to incoming: A) async throws -> Void where A.Element == String {
         for try await message in incoming {
-            try await onMessage(message, self)
-
             // Detect and ignore error responses.
             if message.starts(with: "44") {
                 // TODO: Determine what to do with returned error messages

Sources/GraphQLWS/Server.swift

@@ -17,7 +17,6 @@ public actor Server<
     let onExecute: (GraphQLRequest, InitPayloadResult) async throws -> GraphQLResult
     let onSubscribe: (GraphQLRequest, InitPayloadResult) async throws -> SubscriptionSequenceType
 
-    let onMessage: (String) async throws -> Void
     let onOperationComplete: (String) async throws -> Void
     let onOperationError: (String, [Error]) async throws -> Void
 
@@ -34,23 +33,20 @@ public actor Server<
     ///   - messenger: The messenger to bind the server to.
     ///   - onExecute: Callback run during `start` resolution for non-streaming queries. Typically this is `API.execute`.
     ///   - onSubscribe: Callback run during `start` resolution for streaming queries. Typically this is `API.subscribe`.
-    ///   - onMessage: Optional callback run on every message event
     ///   - onOperationComplete: Optional callback run when an operation completes
     ///   - onOperationError: Optional callback run when an operation errors
     public init(
         messenger: Messenger,
         onInit: @escaping (InitPayload) async throws -> InitPayloadResult,
         onExecute: @escaping (GraphQLRequest, InitPayloadResult) async throws -> GraphQLResult,
         onSubscribe: @escaping (GraphQLRequest, InitPayloadResult) async throws -> SubscriptionSequenceType,
-        onMessage: @escaping (String) async throws -> Void = { _ in },
         onOperationComplete: @escaping (String) async throws -> Void = { _ in },
         onOperationError: @escaping (String, [Error]) async throws -> Void = { _, _ in }
     ) {
         self.messenger = messenger
         self.onInit = onInit
         self.onExecute = onExecute
         self.onSubscribe = onSubscribe
-        self.onMessage = onMessage
         self.onOperationComplete = onOperationComplete
         self.onOperationError = onOperationError
     }
@@ -59,8 +55,6 @@ public actor Server<
     /// - Parameter incoming: The client message sequence that the server should react to.
     public func listen<A: AsyncSequence & Sendable>(to incoming: A) async throws -> Void where A.Element == String {
         for try await message in incoming {
-            try await onMessage(message)
-
             // Detect and ignore error responses.
             if message.starts(with: "44") {
                 // TODO: Determine what to do with returned error messages

Tests/GraphQLWSTests/GraphQLWSTests.swift

@@ -29,24 +29,23 @@ struct GraphqlTransportWSTests {
             }
         )
         let (messageStream, messageContinuation) = AsyncThrowingStream<String, any Error>.makeStream()
+        let serverMessageStream = serverMessenger.stream.map { message in
+            messageContinuation.yield(message)
+            // Expect only one message
+            messageContinuation.finish()
+            return message
+        }
         let client = Client<TokenInitPayload>(
             messenger: clientMessenger,
             onError: { message, _ in
                 messageContinuation.finish(throwing: message.payload[0])
-            },
-            onMessage: { message, _ in
-                messageContinuation.yield(message)
-                // Expect only one message
-                messageContinuation.finish()
+                await clientMessenger.close()
             }
         )
-        let serverStream = serverMessenger.stream
         let clientStream = clientMessenger.stream
         Task {
             try await server.listen(to: clientStream)
-        }
-        Task {
-            try await client.listen(to: serverStream)
+            await serverMessenger.close()
         }
 
         try await client.sendStart(
@@ -59,6 +58,7 @@ struct GraphqlTransportWSTests {
             ),
             id: UUID().uuidString
         )
+        try await client.listen(to: serverMessageStream)
 
         let messages = try await messageStream.reduce(into: [String]()) { result, message in
             result.append(message)
@@ -92,31 +92,31 @@ struct GraphqlTransportWSTests {
             }
         )
         let (messageStream, messageContinuation) = AsyncThrowingStream<String, any Error>.makeStream()
+        let serverMessageStream = serverMessenger.stream.map { message in
+            messageContinuation.yield(message)
+            // Expect only one message
+            messageContinuation.finish()
+            return message
+        }
         let client = Client<TokenInitPayload>(
             messenger: clientMessenger,
             onError: { message, _ in
                 messageContinuation.finish(throwing: message.payload[0])
-            },
-            onMessage: { message, _ in
-                messageContinuation.yield(message)
-                // Expect only one message
-                messageContinuation.finish()
+                await clientMessenger.close()
             }
         )
-        let serverStream = serverMessenger.stream
         let clientStream = clientMessenger.stream
         Task {
             try await server.listen(to: clientStream)
-        }
-        Task {
-            try await client.listen(to: serverStream)
+            await serverMessenger.close()
         }
 
         try await client.sendConnectionInit(
             payload: TokenInitPayload(
                 authToken: ""
             )
         )
+        try await client.listen(to: serverMessageStream)
 
         let messages = try await messageStream.reduce(into: [String]()) { result, message in
             result.append(message)
@@ -150,6 +150,10 @@ struct GraphqlTransportWSTests {
             }
         )
         let (messageStream, messageContinuation) = AsyncThrowingStream<String, any Error>.makeStream()
+        let serverMessageStream = serverMessenger.stream.map { message in
+            messageContinuation.yield(message)
+            return message
+        }
         let client = Client<TokenInitPayload>(
             messenger: clientMessenger,
             onConnectionAck: { _, client in
@@ -166,24 +170,21 @@ struct GraphqlTransportWSTests {
             },
             onError: { message, _ in
                 messageContinuation.finish(throwing: message.payload[0])
+                await clientMessenger.close()
             },
             onComplete: { _, _ in
                 messageContinuation.finish()
-            },
-            onMessage: { message, _ in
-                messageContinuation.yield(message)
+                await clientMessenger.close()
             }
         )
-        let serverStream = serverMessenger.stream
         let clientStream = clientMessenger.stream
         Task {
             try await server.listen(to: clientStream)
-        }
-        Task {
-            try await client.listen(to: serverStream)
+            await serverMessenger.close()
         }
 
         try await client.sendConnectionInit(payload: TokenInitPayload(authToken: ""))
+        try await client.listen(to: serverMessageStream)
 
         let messages = try await messageStream.reduce(into: [String]()) { result, message in
             result.append(message)
@@ -222,6 +223,11 @@ struct GraphqlTransportWSTests {
             }
         )
         let (messageStream, messageContinuation) = AsyncThrowingStream<String, any Error>.makeStream()
+        // Used to extract the server messages
+        let serverMessageStream = serverMessenger.stream.map { message in
+            messageContinuation.yield(message)
+            return message
+        }
         let client = Client<TokenInitPayload>(
             messenger: clientMessenger,
             onConnectionAck: { _, client in
@@ -250,24 +256,21 @@ struct GraphqlTransportWSTests {
             },
             onError: { message, _ in
                 messageContinuation.finish(throwing: message.payload[0])
+                await clientMessenger.close()
             },
             onComplete: { _, _ in
                 messageContinuation.finish()
-            },
-            onMessage: { message, _ in
-                messageContinuation.yield(message)
+                await clientMessenger.close()
             }
         )
         let clientStream = clientMessenger.stream
-        let serverStream = serverMessenger.stream
         Task {
             try await server.listen(to: clientStream)
-        }
-        Task {
-            try await client.listen(to: serverStream)
+            await serverMessenger.close()
         }
 
         try await client.sendConnectionInit(payload: TokenInitPayload(authToken: ""))
+        try await client.listen(to: serverMessageStream)
 
         let messages = try await messageStream.reduce(into: [String]()) { result, message in
             result.append(message)