Swift Concurrency + ViewModel action handling

Effect.swift

enum Effect {
  /// Action completed synchronously
  case none

  /// Action produced a task
  case task(Task<Void, Never>)

  /// Notify about completion in a closure
  func onCompletion(_ completion: @escaping () -> Void) {
    switch self {
    case .none:
      completion()

    case let .task(task):
      Task {
        _ = await task.value
        completion()
      }
    }
  }

  /// Wait until completion
  var completion: Void {
    get async {
      switch self {
      case .none:
        return

      case let .task(task):
        await withTaskCancellationHandler {
          _ = await task.value
        } onCancel: {
          task.cancel()
        }
      }
    }
  }
}

Sample.swift

import SwiftUI

final class ViewModel: ObservableObject {
  enum Action {
    case buttonTap
  }

  private let tasks = TaskEffectHandler()

  func connect() async {
    await tasks.connect()
  }

  @discardableResult
  func handle(_ action: Action) -> Effect {
    switch action {
    case .buttonTap:
      return tasks.add {
        do {
          try await Task.sleep(for: .seconds(2))
        } catch {
          print("Task cancelled")
        }
      }
    }
  }
}

struct ContentView: View {
  @ObservedObject var vm = ViewModel()
  
  @State var flag = true

  var body: some View {
    VStack {
      Toggle("Show", isOn: $flag)

      Spacer()

      if flag {
        AsyncButton {
          await vm.handle(.buttonTap).completion
        } label: { isLoading in
          if isLoading {
            ProgressView()
          } else {
            Text("Tap me")
          }
        }
        .task { await vm.connect() }
        .frame(maxWidth: .infinity, maxHeight: .infinity)
      }
    }
    .padding()
  }
}

struct AsyncButton<Label: View>: View {
  let action: @Sendable () async -> Void

  let label: (_ isLoading: Bool) -> Label

  @State private var id: UUID?

  @State private var isLoading = false

  init(
    @_inheritActorContext action: @Sendable @escaping () async -> Void,
    @ViewBuilder label: @escaping (_ isLoading: Bool) -> Label
  ) {
    self.action = action
    self.label = label
  }

  var body: some View {
    Button {
      id = .init()
    } label: {
      label(isLoading)
    }
    .disabled(isLoading)
    .task(id: id) {
      guard id != nil else { return }
      isLoading = true
      await action()
      guard !Task.isCancelled else { return }
      isLoading = false
    }

  }
}

struct ContentView_Previews: PreviewProvider {
  static var previews: some View {
    ContentView()
  }
}

TaskEffectHandler.swift

final class TaskEffectHandler {
  private typealias OperationStream = AsyncStream<@Sendable () async -> Void>

  private var operationStream: OperationStream?

  private var operationStreamContinuation: OperationStream.Continuation?

  func connect() async {
    operationStreamContinuation?.finish()
    let stream = OperationStream.makeStream() // Swift 5.9
    (operationStream, operationStreamContinuation) = stream
    await withDiscardingTaskGroup { group in // Swift 5.9
      for await operation in stream.stream {
        group.addTask {
          await operation()
        }
      }
    }
    (operationStream, operationStreamContinuation) = (nil, nil)
  }

  func add(_ operation: @escaping @Sendable () async -> Void) -> Effect {
    guard let operationStreamContinuation else {
      print("⚠️ Warning: adding operations while not connected")
      return .task(
        Task {
          await operation()
        }
      )
    }

    let (cancelStream, cancelContinuation) = AsyncStream<Void>.makeStream()

    return .task(
      Task {
        await withTaskCancellationHandler {
          await withCheckedContinuation { continuation in
            operationStreamContinuation.yield {
              let operationTask = Task {
                await operation()
                cancelContinuation.finish()
              }

              for await _ in cancelStream {
                operationTask.cancel()
              }

              if Task.isCancelled {
                operationTask.cancel()
              }

              continuation.resume()
            }
          }
        } onCancel: {
          cancelContinuation.yield()
        }
      }
    )
  }
}