What would be the best way to convert a system callback like adding a NotificationCenter.default.addObserver and also removing the listener on error / cancellation wrapped in an async function?

The usage I strife for is:

@MainActor func waitForLifecycleEvent(_ applicationState: UIApplication.State = .active) async 

In kotlin there is a [suspendCancellableCoroutine](https://kotlinlang.org/api/kotlinx.coroutines/kotlinx-coroutines-
core/kotlinx.coroutines/suspend-cancellable-coroutine.html)

So that you could implement it like this: https://gist.github.com/kibotu/bd4eb5d37410a18311d200103a827df8

Where you can use add a callback but also have a way to properly cleanup again.

In swift we have withTaskCancellationHandler which can deal with cancelation but it does not support a callback since it’s operation returns immediately.

And we have withCheckedThrowingContinuation which does handle callbacks but does not handle cancellation which means we have no way of cleaning up again.

And even combining those two would not be enough due to the way I variables are being shared between tasks so that for all cases to cleanup we would additionally need a wrapper object with deinit to cleanup.

My first crude seemingly working approach looks like this:

import UIKit


@MainActor
public func waitForLifecycleEvent(_ applicationState: UIApplication.State = .active) async {
    // Return immediately if already in foreground
    if UIApplication.shared.applicationState == applicationState { return }

    let wrapper = ObserverWrapper()

    do {
        try await withTaskCancellationHandler(
            operation: {
            // Wait for notification asynchronously.
            try await withCheckedThrowingContinuation { continuation in
                switch applicationState {
                case .active:
                    // Register for the notification.
                    let observer = NotificationCenter.default.addObserver(forName: UIApplication.willEnterForegroundNotification, object: nil, queue: .main) { _ in
                        continuation.resume()
                    }
                    wrapper.observer = observer

                case .inactive, .background:
                    // Register for the notification.
                    let observer = NotificationCenter.default.addObserver(forName: UIApplication.didEnterBackgroundNotification, object: nil, queue: .main) { _ in
                        continuation.resume()
                    }
                    wrapper.observer = observer
                @unknown default:
                    CorePlugin.logInfo("[waitForLifecycleEvent] unknown default (applicationState)")
                }
            }
            }, onCancel: {
            // Cancellation handler is called if the task is cancelled. We need to unregister observer.
            wrapper.observer.map(NotificationCenter.default.removeObserver)
            })
    } catch {
        // ignore
    }
}

enum AwaitForForegroundError: Error {
    case cancelled
}

private final class ObserverWrapper {
    var observer: NSObjectProtocol?
    deinit {
        if let observer = observer {
            NotificationCenter.default.removeObserver(observer)
        }
    }
}

This all looks like a bit much considering. Is there a better way to? Or are there any improvements or gotchas you can recommend?

Thanks in advance