diff --git a/Sources/AsyncWebSocketClient/AsyncWebSocketClient.swift b/Sources/AsyncWebSocketClient/AsyncWebSocketClient.swift
index 0632e7a..a8d1ea6 100644
--- a/Sources/AsyncWebSocketClient/AsyncWebSocketClient.swift
+++ b/Sources/AsyncWebSocketClient/AsyncWebSocketClient.swift
@@ -135,8 +135,8 @@ public actor AsyncWebSocketClient: NSObject, AsyncWebSocketClientProtocol {
   /// A debouncing behavior is implemented to send a ping-pong every time after 20 seconds
   /// have elapsed since the last time an event happened or this function was initially called.
   func startPingPongHandler() {
-    Task { [weak self] in
-      guard let stream = await self?.streamGenerator.subscribe() else { return }
+    Task { [weak self, streamGenerator, scheduler] in
+      let stream = await streamGenerator.subscribe()
 
       let debouncedStream = stream.debounce(for: Constants.debounceTime, scheduler: scheduler)
 
diff --git a/Sources/AsyncWebSocketClient/StreamGenerator/StreamGenerator.swift b/Sources/AsyncWebSocketClient/StreamGenerator/StreamGenerator.swift
index 9e014b0..69bf753 100644
--- a/Sources/AsyncWebSocketClient/StreamGenerator/StreamGenerator.swift
+++ b/Sources/AsyncWebSocketClient/StreamGenerator/StreamGenerator.swift
@@ -7,8 +7,6 @@
 
 import Foundation
 
-// OPTIONAL TODO: move to a different package
-
 actor StreamGenerator<T> {
   var subscribers = [UUID: AsyncStream<T>.Continuation]()
 
@@ -48,18 +46,7 @@ actor StreamGenerator<T> {
     for key in subscribers.keys {
       guard let subscriber = subscribers[key] else { continue }
       subscriber.finish()
-      removeSubscriber(with: key)
+      subscribers.removeValue(forKey: key)
     }
   }
 }
-
-//public protocol WriteStreamGenerator: Actor {
-//    associatedtype Value
-//    func updateValue(_ value: Value)
-//}
-//
-//public protocol ReadStreamGenerator: Actor {
-//    associatedtype Value
-//    var value: Value { get }
-//    func subscribe() -> AsyncStream<Value>
-//}